# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# pylint: disable=invalid-name, unused-variable, too-many-locals, unused-argument
# ruff: noqa: F841
"""Depthwise convolution operators"""

from collections import namedtuple

import numpy as np

import tvm
from tvm import te

from ..utils import get_const_tuple, simplify
from .dilate import dilate
from .pad import pad
from .utils import get_pad_tuple

# workload description of depthwise-conv2d
Workload = namedtuple(
    "Workload",
    [
        "in_dtype",
        "out_dtype",
        "height",
        "width",
        "in_filter",
        "out_filter",
        "kernel_h",
        "kernel_w",
        "padt",
        "padl",
        "padb",
        "padr",
        "dilation_h",
        "dilation_w",
        "stride_h",
        "stride_w",
    ],
)


def _get_workload(data, kernel, stride, padding, dilation, out_dtype, data_layout="NCHW"):
    """Get the workload structure for a depthwise conv2d.

    Input data and filter should use NCHW layout.
    """
    if data_layout == "NCHW":
        _, in_channel, height, width = get_const_tuple(data.shape)
        filter_channel, channel_multiplier, kh, kw = get_const_tuple(kernel.shape)
    elif data_layout == "NHWC":
        _, height, width, in_channel = get_const_tuple(data.shape)
        kh, kw, filter_channel, channel_multiplier = get_const_tuple(kernel.shape)
    elif data_layout == "NCHWc":
        _, in_channel_chunk, height, width, in_channel_block = get_const_tuple(data.shape)
        in_channel = in_channel_chunk * in_channel_block
        (filter_channel_chunk, cm_chunk, kh, kw, cm_block, filter_channel_block) = get_const_tuple(
            kernel.shape
        )
        filter_channel = filter_channel_chunk * filter_channel_block
        channel_multiplier = cm_chunk * cm_block

        assert in_channel_block == filter_channel_block, (
            f"Incorrect dimensions, data has block size {in_channel_block}, but filter has "
            f"block size {filter_channel_block}"
        )

    else:
        raise ValueError(f"Data layout {data_layout} not supported")

    assert in_channel == filter_channel, (
        f"Incorrect dimensions, data has {in_channel} channels but filter expects "
        f"{filter_channel} channels"
    )

    out_channel = filter_channel * channel_multiplier
    dilation_h, dilation_w = (
        dilation if isinstance(dilation, tuple | list) else (dilation, dilation)
    )
    if isinstance(stride, tuple | list):
        HSTR, WSTR = stride
    else:
        HSTR, WSTR = stride, stride
    assert (data.dtype == kernel.dtype) or (data.dtype == "uint8" and kernel.dtype == "int8"), (
        f"Do not support inputs with different data types now. {data.dtype} vs. {kernel.dtype}"
    )
    dilated_kernel_h = (kh - 1) * dilation_h + 1
    dilated_kernel_w = (kw - 1) * dilation_w + 1
    pt, pl, pb, pr = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
    return Workload(
        data.dtype,
        out_dtype,
        height,
        width,
        in_channel,
        out_channel,
        kh,
        kw,
        pt,
        pl,
        pb,
        pr,
        dilation_h,
        dilation_w,
        HSTR,
        WSTR,
    )


def depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=None):
    """Depthwise convolution nchw forward operator.

    Parameters
    ----------
    Input : tvm.te.Tensor
        4-D with shape [batch, in_channel, in_height, in_width]

    Filter : tvm.te.Tensor
        4-D with shape [in_channel, channel_multiplier, filter_height, filter_width]

    stride : int or a list/tuple of two ints
        The spatial stride, or (stride_height, stride_width).

    padding : int or str
        Padding size, or ['VALID', 'SAME']

    dilation: int or a list/tuple of two ints
        dilation size, or [dilation_height, dilation_width]

    out_dtype: str, optional
        Output data type

    Returns
    -------
    Output : tvm.te.Tensor
        4-D with shape [batch, out_channel, out_height, out_width]
    """
    out_dtype = Input.dtype if out_dtype is None else out_dtype

    if isinstance(stride, int):
        stride_h = stride_w = stride
    else:
        stride_h, stride_w = stride

    if isinstance(dilation, int):
        dilation_h = dilation_w = dilation
    else:
        dilation_h, dilation_w = dilation

    batch, in_channel, in_height, in_width = Input.shape
    # shape of dilated kernel
    filter_channel, channel_multiplier, filter_height, filter_width = Filter.shape

    dilated_kernel_h = (filter_height - 1) * dilation_h + 1
    dilated_kernel_w = (filter_width - 1) * dilation_w + 1
    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
        padding, (dilated_kernel_h, dilated_kernel_w)
    )
    out_channel = simplify(in_channel * channel_multiplier)
    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)

    # padding stage
    pad_before = [0, 0, pad_top, pad_left]
    pad_after = [0, 0, pad_down, pad_right]
    PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
    # depthconv stage
    idxdiv = tvm.tirx.indexdiv
    idxmod = tvm.tirx.indexmod
    di = te.reduce_axis((0, filter_height), name="di")
    dj = te.reduce_axis((0, filter_width), name="dj")
    Output = te.compute(
        (batch, out_channel, out_height, out_width),
        lambda b, c, i, j: te.sum(
            (
                PaddedInput[
                    b,
                    idxdiv(c, channel_multiplier),
                    i * stride_h + di * dilation_h,
                    j * stride_w + dj * dilation_w,
                ].astype(out_dtype)
                * Filter[
                    idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier), di, dj
                ].astype(out_dtype)
            ),
            axis=[di, dj],
        ),
        name="DepthwiseConv2d",
        tag="depthwise_conv2d_nchw",
    )
    return Output


def depthwise_conv2d_nhwc(
    Input, Filter, stride, padding, dilation, kernel_layout="HWOI", out_dtype=None
):
    """Depthwise convolution nhwc forward operator.

    Parameters
    ----------
    Input : tvm.te.Tensor
        4-D with shape [batch, in_height, in_width, in_channel]

    Filter : tvm.te.Tensor
        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]

    stride : tuple of two ints
        The spatial stride along height and width

    padding : int or str
        Padding size, or ['VALID', 'SAME']

    dilation: int or a list/tuple of two ints
        dilation size, or [dilation_height, dilation_width]

    out_dtype: str, optional
        Output data type

    Returns
    -------
    Output : tvm.te.Tensor
        4-D with shape [batch, out_height, out_width, out_channel]
    """
    out_dtype = Input.dtype if out_dtype is None else out_dtype

    if isinstance(stride, int):
        stride_h = stride_w = stride
    else:
        stride_h, stride_w = stride

    if isinstance(dilation, int):
        dilation_h = dilation_w = dilation
    else:
        dilation_h, dilation_w = dilation

    batch, in_height, in_width, in_channel = Input.shape

    # shape of dilated kernel
    if kernel_layout == "HWIO":
        filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape
        kernel_permutation = [0, 1, 3, 2]
    else:
        filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape
        kernel_permutation = [0, 1, 2, 3]

    dilated_kernel_h = (filter_height - 1) * dilation_h + 1
    dilated_kernel_w = (filter_width - 1) * dilation_w + 1
    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
        padding, (dilated_kernel_h, dilated_kernel_w)
    )
    out_channel = simplify(in_channel * channel_multiplier)
    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)

    # padding stage
    pad_before = [0, pad_top, pad_left, 0]
    pad_after = [0, pad_down, pad_right, 0]
    PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
    # depthconv stage
    idxdiv = tvm.tirx.indexdiv
    idxmod = tvm.tirx.indexmod

    di = te.reduce_axis((0, filter_height), name="di")
    dj = te.reduce_axis((0, filter_width), name="dj")
    Output = te.compute(
        (batch, out_height, out_width, out_channel),
        lambda b, i, j, c: te.sum(
            (
                PaddedInput[
                    b,
                    i * stride_h + di * dilation_h,
                    j * stride_w + dj * dilation_w,
                    idxdiv(c, channel_multiplier),
                ].astype(out_dtype)
                * Filter[
                    tuple(
                        np.array(
                            [di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)]
                        )[kernel_permutation]
                    )
                ].astype(out_dtype)
            ),
            axis=[di, dj],
        ),
        name="DepthwiseConv2d",
        tag="depthwise_conv2d_nhwc",
    )
    return Output


def depthwise_conv2d_backward_input_nhwc(Filter, Out_grad, oshape, ishape, stride, padding):
    """Depthwise convolution nhwc backward wrt input operator.

    Parameters
    ----------
    Filter : tvm.te.Tensor
        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]

    Out_grad : tvm.te.Tensor
        4-D with shape [batch, out_height, out_width, out_channel]

    stride : tuple of two ints
        The spatial stride along height and width

    padding : int or str
        Padding size, or ['VALID', 'SAME']

    Returns
    -------
    Output : tvm.te.Tensor
        4-D with shape [batch, in_height, in_width, in_channel]
    """
    batch, in_h, in_w, in_c = ishape
    _, out_h, out_w, out_c = oshape
    filter_h, filter_w, _, channel_multiplier = Filter.shape
    if isinstance(stride, int):
        stride_h = stride_w = stride
    else:
        stride_h, stride_w = stride

    dilated_out_grad = dilate(Out_grad, [1, stride_h, stride_w, 1], name="dilated_out_grad")

    # padding params in forward propagation
    fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
    # padding params in backward propagation
    bpad_top = filter_h - 1 - fpad_top
    bpad_bottom = (filter_h - 1 - fpad_bottom) + (stride_h - 1)
    bpad_left = filter_w - 1 - fpad_left
    bpad_right = (filter_w - 1 - fpad_right) + (stride_w - 1)

    padded_out_grad = pad(
        dilated_out_grad,
        [0, bpad_top, bpad_left, 0],
        [0, bpad_bottom, bpad_right, 0],
        name="padded_out_grad",
    )

    dh = te.reduce_axis((0, filter_h), name="dh")
    dw = te.reduce_axis((0, filter_w), name="dw")
    dc = te.reduce_axis((0, channel_multiplier), name="dc")

    In_grad = te.compute(
        (batch, in_h, in_w, in_c),
        lambda b, h, w, c: te.sum(
            padded_out_grad[b, h + dh, w + dw, c * channel_multiplier + dc]
            * Filter[filter_h - 1 - dh, filter_w - 1 - dw, c, dc],
            axis=[dh, dw, dc],
        ),
        tag="depthwise_conv2d_backward_input_nhwc",
    )

    return In_grad


def depthwise_conv2d_backward_weight_nhwc(Input, Out_grad, oshape, fshape, stride, padding):
    """Depthwise convolution nhwc backward wrt weight operator.

    Parameters
    ----------
    Input : tvm.te.Tensor
        4-D with shape [batch, in_height, in_width, in_channel]

    Out_grad : tvm.te.Tensor
        4-D with shape [batch, out_height, out_width, out_channel]

    stride : tuple of two ints
        The spatial stride along height and width

    padding : int or str
        Padding size, or ['VALID', 'SAME']

    Returns
    -------
    Output : tvm.te.Tensor
        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
    """
    batch, out_h, out_w, out_c = oshape
    filter_h, filter_w, _, channel_multiplier = fshape
    in_c = Input.shape[3].value
    if isinstance(stride, int):
        stride_h = stride_w = stride
    else:
        stride_h, stride_w = stride

    pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (filter_h, filter_w))

    padded_in = pad(
        Input, [0, pad_top, pad_left, 0], [0, pad_bottom, pad_right, 0], name="padded_in"
    )

    dh = te.reduce_axis((0, Out_grad.shape[1].value), name="dh")
    dw = te.reduce_axis((0, Out_grad.shape[2].value), name="dw")
    db = te.reduce_axis((0, batch), name="db")
    idxdiv = tvm.tirx.indexdiv
    idxmod = tvm.tirx.indexmod

    Weight_grad = te.compute(
        (filter_h, filter_w, in_c, channel_multiplier),
        lambda fh, fw, c, m: te.sum(
            Out_grad[db, dh, dw, c * channel_multiplier + idxmod(m, channel_multiplier)]
            * padded_in[db, fh + dh * stride_h, fw + dw * stride_w, c],
            axis=[db, dh, dw],
        ),
        tag="depthwise_conv2d_backward_weight_nhwc",
    )

    return Weight_grad


def depthwise_conv2d_NCHWc(
    Input, Filter, stride, padding, dilation, layout, out_layout, out_dtype=None
):
    """Depthwise convolution NCHW[x]c forward operator.

    Parameters
    ----------
    Input : tvm.te.Tensor
        5-D with shape [batch, in_channel_chunk, in_height, in_width, in_channel_block]

    Filter : tvm.te.Tensor
        6-D with shape [out_channel_chunk, 1, filter_height, filter_width, 1, out_channel_block]
        In NCHWc depthwise convolution,
        we group kernel's in_channel and channel_multiplier together then do the tiling.

    stride : tuple of two ints
        The spatial stride along height and width

    padding : int or str
        Padding size, or ['VALID', 'SAME']

    dilation: int or a list/tuple of two ints
         dilation size, or [dilation_height, dilation_width]

    layout : str
        Input data layout

    out_layout : str
        Output data layout

    out_dtype: str, optional
        Output data type

    Returns
    -------
    Output : tvm.te.Tensor
        5-D with shape [batch, out_channel_chunk, out_height, out_width, out_channel_block]
    """
    raise ValueError("missing register for topi.nn.depthwise_conv2d_NCHWc")