#!/usr/bin/env python3
# coding: utf-8
# Copyright 2019 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""operator dsl function: maxpool_grad"""
import akg.tvm
from akg.tvm.hybrid import script
from akg.utils import custom_tiling as ct_util, kernel_exec as utils, \
    validation_check as vc_util
from akg.utils.dsl_create import cal_pad_shapes_by_strategy
from akg.dim import DIM
from akg.utils.format_transform import get_shape


def get_attrs():
    """get attrs config"""
    attr_map = {
        "disable_cse": 1,
        "pragma_reschedule": 1,
        "pragma_disable_schedule_shift": 1,
        "pragma_opt_for_davinci": 1,
        "pragma_disable_loop_reversal": 1,
        "loop_partition_unroll": False,
        "enable_pre_poly_loop_partition": False
    }
    return attr_map


set_attr_map_ = {
    str(((2, 16, 40, 24, 16), (1, 1), (2, 2), (0, 0, 0, 0))): (
        ("pragma_reschedule", 0),),
}

maxpool_grad_dim_map = {
    str(((2, 16, 40, 24, 16), (1, 1), (2, 2), (0, 0, 0, 0))): (
        (1, 1), (2, 1), (20, 1), (16, 1)),
    str(((32, 4, 112, 112, 16), (3, 3), (2, 2), 'SAME')): (
        (1, 1), (1, 1), (16, 1), (1, 1)),
    str(((32, 4, 112, 112, 16), (3, 3), (2, 2), (0, 1, 0, 1))): (
        (1, 1), (1, 1), (16, 1), (1, 1)),
    str(((32, 4, 112, 112, 16), (3, 3), (2, 2), (1, 0, 1, 0))): (
        (1, 1), (1, 1), (16, 1), (1, 1)),
    str(((1, 1, 4, 4, 16), (2, 2), (2, 2), (0, 0, 0, 0))): (
        (0, 1), (0, 1), (0, 1)),
    str(((1, 1, 16, 16, 16), (4, 4), (4, 4), (0, 0, 0, 0))): (
        (0, 1), (0, 1), (0, 1)),
    str(((1, 1, 32, 32, 16), (4, 4), (4, 4), (0, 0, 0, 0))): (
        (0, 1), (0, 1), (0, 1)),
    str(((32, 16, 28, 28, 16), (3, 3), (2, 2), 'VALID')): (
        (8, 1), (1, 1), (16, 1), (1, 1)),
    str(((32, 16, 13, 13, 16), (3, 3), (2, 2), 'VALID')): (
        (8, 1), (1, 1), (16, 1), (1, 1)),
    str(((32, 6, 57, 57, 16), (3, 3), (2, 2), 'VALID')): (
        (4, 1), (1, 1), (16, 1), (1, 1)),
}


def maxpool_grad_set_dim_func(x, y, dy, kernel, stride, pad):
    """dim func for maxpool grad"""
    key = str((tuple(x.shape), tuple(kernel), tuple(stride), pad))
    attrs = {}
    if key in set_attr_map_.keys():
        for attr in set_attr_map_[key]:
            attrs[attr[0]] = attr[1]
    if key in maxpool_grad_dim_map.keys():
        return ct_util.set_dims_by_key(key, maxpool_grad_dim_map), key, attrs
    return "", key, attrs


@vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor,
                          akg.tvm.tensor.Tensor, (list, tuple), (list, tuple),
                          (str, list, tuple))
def maxpool_grad(x, y, dy, kernel, stride, pad):
    """
    Performs the gradient of maxpool pooling on the input datas.

    Note:
        Only support 5D format(NC1HWC0), and pooling will work on H and W.

    Args:
        x (tvm.tensor.Tensor): Tensor of type float16, float32.
        y (tvm.tensor.Tensor): Tensor, the maxpool result.
        dy (tvm.tensor.Tensor): Tensor, the gradient needed to be propagation.
        kernel (Union[List, Tuple]): two int numbers for pooling window's size.
        stride (Union[List, Tuple]): two int numbers for window's stride.
        pad (Union[String, List, Tuple]): padding, should be 'VALID','SAME' or
            instance of list(four int numbers, as 'CONSTANTS' strategy).
            Support **pad** is the same as avgpool's **Strategies**.

    Returns:
        Tensor as result for gradient of maxpooling.
    """
    attrs = get_attrs()
    dim_info, _, attrs_info = maxpool_grad_set_dim_func(x, y, dy, kernel, stride, pad)
    attrs.update(attrs_info)
    attrs[DIM] = dim_info

    shape = get_shape(x)
    ori_dtype = x.dtype
    vc_util.ops_dtype_check(ori_dtype, vc_util.DtypeForDavinci.ALL_FLOAT)

    if utils.product_is_mini() and ori_dtype == 'float32':
        raise RuntimeError("Maxpool only support"
                           "\'float16\' while platform is mini_v100!")
    dtype = ori_dtype

    if len(shape) != 5:
        raise ValueError("Only support 5-dim pooling!")
    if shape[-1] % 16 != 0:
        raise ValueError("Last shape must be divisible by 16!")
    if len(kernel) != 2:
        raise ValueError("Only support 2-dim kernel!")
    if len(stride) != 2:
        raise ValueError("Only support 2-dim stride!")
    if not isinstance(pad, str) \
            and not (isinstance(pad, (list, tuple)) and len(pad) == 4):
        raise ValueError("Only support string or list/tuple of 4 int numbers!")

    vc_util.check_shape(shape)

    in_n, in_c1, in_h, in_w, in_c0 = shape
    k_h, k_w = kernel
    s_h, s_w = stride
    [ph_h, ph_t, pw_h, pw_t], [y_h, y_w] = \
        cal_pad_shapes_by_strategy(shape, kernel, stride, pad)
    k_h_hybrid = k_h
    k_w_hybrid = k_w

    yn = in_n
    yc1 = in_c1
    yc0 = in_c0

    @script(capture=locals())
    def max_pool_grad_hybrid(zero_, one_, min_value_, x_, y_, dy_):
        x_dummy_ = allocate((in_n, in_c1, ph_h + in_h + ph_t, pw_h + in_w + pw_t, in_c0),
                            x_.dtype, "local")
        x_img_ = allocate((yn, yc1, y_h, y_w, k_h_hybrid, k_w_hybrid, yc0),
                          x_.dtype, "local")
        y_img_ = allocate((yn, yc1, y_h, y_w, k_h_hybrid, k_w_hybrid, yc0),
                          x_.dtype)
        mask_ = allocate((yn, yc1, y_h, y_w, k_h_hybrid, k_w_hybrid, yc0),
                         x_.dtype)
        mask_new = allocate((yn, yc1, y_h, y_w, k_h_hybrid, k_w_hybrid, yc0),
                            dy_.dtype)
        mask_res = allocate((yn, yc1, y_h, y_w, k_h_hybrid, k_w_hybrid, yc0),
                            dy_.dtype)
        output_pre = allocate((yn, yc1, y_h, y_w, k_h_hybrid, k_w_hybrid, yc0),
                              dy_.dtype)
        output_dummy_body = allocate((in_n, in_c1,
                                      ph_h + in_h + ph_t, pw_h + in_w + pw_t, in_c0), dy_.dtype)
        output = output_tensor((in_n, in_c1, in_h, in_w, in_c0), dy_.dtype)

        for n in range(yn):
            for c1 in range(yc1):
                for h in range(y_h):

                    for kh in range(k_h_hybrid):
                        for iw in range(pw_h + in_w + pw_t):
                            for c0 in range(yc0):
                                x_dummy_[n, c1,
                                         h * s_h + kh, iw, c0] = min_value_
                                output_dummy_body[n, c1,
                                                  h * s_h + kh, iw, c0] = zero_

                    for kh in range(k_h_hybrid):
                        for iw in range(in_w):
                            for c0 in range(yc0):
                                if (h * s_h + kh >= ph_h
                                        and h * s_h + kh < in_h + ph_h):
                                    x_dummy_[n, c1, h * s_h + kh,
                                             iw + pw_h, c0] = \
                                        x_[n, c1, h * s_h + kh - ph_h, iw, c0]

                    for kh in range(k_h_hybrid):
                        for iw in range(in_w):
                            for c0 in range(yc0):
                                if (h * s_h + kh >= ph_h
                                        and h * s_h + kh < in_h + ph_h):
                                    output_dummy_body[n, c1,
                                                      h * s_h + kh, iw + pw_h, c0] = \
                                        output[n, c1, h * s_h + kh - ph_h, iw, c0]

                    for w in range(y_w):
                        for kh in range(k_h_hybrid):
                            for kw in range(k_w_hybrid):
                                for c0 in range(yc0):
                                    x_img_[n, c1, h, w, kh, kw, c0] = \
                                        x_dummy_[n, c1, h * s_h + kh,
                                                 w * s_w + kw, c0]
                                    y_img_[n, c1, h, w, kh, kw, c0] = \
                                        y_[n, c1, h, w, c0]
                                    mask_[n, c1, h, w, kh, kw, c0] = zero_ \
                                        if x_img_[n, c1, h, w, kh, kw, c0] \
                                        < y_img_[n, c1, h, w, kh, kw, c0] \
                                        else one_
                        for kh in range(k_h_hybrid):
                            for kw in range(k_w_hybrid):
                                for c0 in range(yc0):
                                    mask_new[n, c1, h, w, kh,
                                             kw, c0] = zero_
                                for kh_0 in range(kh):
                                    for kw_0 in range(k_w_hybrid):
                                        for c0 in range(yc0):
                                            mask_new[n, c1, h, w, kh,
                                                     kw, c0] = \
                                                mask_new[n, c1, h, w,
                                                         kh, kw, c0] \
                                                + mask_[n, c1, h, w,
                                                        kh_0, kw_0, c0]
                                for kw_0 in range(kw + 1):
                                    for c0 in range(yc0):
                                        mask_new[n, c1, h, w, kh, kw, c0] = \
                                            mask_new[n, c1, h, w, kh, kw, c0] \
                                            + mask_[n, c1, h, w, kh, kw_0, c0]
                        for kh in range(k_h_hybrid):
                            for kw in range(k_w_hybrid):
                                for c0 in range(yc0):
                                    mask_res[n, c1, h, w, kh, kw, c0] = \
                                        zero_ \
                                        if mask_new[n, c1, h, w, kh, kw, c0] \
                                        > mask_[n, c1, h, w, kh, kw, c0] \
                                        else mask_[n, c1, h, w, kh, kw, c0]
                                    output_pre[n, c1, h, w, kh, kw, c0] = \
                                        mask_res[n, c1, h, w, kh, kw, c0] \
                                        * dy_[n, c1, h, w, c0]
                                    output_dummy_body[n, c1,
                                                      h * s_h + kh, w * s_w + kw, c0] += \
                                        output_pre[n, c1, h, w, kh, kw, c0]
                    for kh in range(k_h_hybrid):
                        for iw in range(in_w):
                            for c0 in range(yc0):
                                if (h * s_h + kh >= ph_h
                                        and h * s_h + kh < in_h + ph_h):
                                    output[n, c1, h * s_h + kh - ph_h,
                                           iw, c0] = \
                                        output_dummy_body[n, c1,
                                                          h * s_h + kh, iw + pw_h, c0]

        return output

    zero = akg.tvm.const(0.0, dtype=dtype)
    one = akg.tvm.const(1.0, dtype=dtype)
    min_value = akg.tvm.const(-65504.0 if dtype == 'float16'
                              else -340282346638528859811704183484516925440.0, dtype=dtype)
    output = max_pool_grad_hybrid(zero, one, min_value, x, y, dy)
    return output, attrs
