# 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
"""affine_grid and grid_sample operator"""

from tvm import te, tirx


def affine_grid(data, target_shape, align_corners=True):
    """affine_grid operator that generates a 2D or 3D sampling grid.

    This operation is described in https://arxiv.org/pdf/1506.02025.pdf. It generates a uniform
    sampling grid within the target shape and normalizes it to [-1, 1]. The provided affine
    transformation is then applied on the sampling grid.

    Parameters
    ----------
    data : tvm.Tensor
        3-D with shape [batch, 2, 3] for 2D or [batch, 3, 4] for 3D. The affine matrix.

    target_shape: list/tuple of int
        Specifies the output spatial shape (H, W) for 2D or (D, H, W) for 3D.

    align_corners : bool
        If True, normalized coordinates map to corner pixels; if False, to pixel centers
        (the PyTorch / ONNX default).

    Returns
    -------
    Output : tvm.Tensor
        [batch, 2, H, W] for 2D or [batch, 3, D, H, W] for 3D.
    """
    assert len(target_shape) in (2, 3)
    if align_corners:
        assert all(s > 1 for s in target_shape), (
            "target spatial dims should be greater than 1 when align_corners is True"
        )

    dtype = data.dtype
    if align_corners:
        starts = [tirx.const(-1.0, dtype=dtype) for _ in target_shape]
        steps = [tirx.const((2.0 - 1e-7) / (s - 1), dtype=dtype) for s in target_shape]
    else:
        # Pixel centers: coordinate i maps to (2 * i + 1) / size - 1.
        starts = [tirx.const(-1.0 + 1.0 / s, dtype=dtype) for s in target_shape]
        steps = [tirx.const(2.0 / s, dtype=dtype) for s in target_shape]

    ndim = len(target_shape)

    def _compute(n, dim, *coords):
        # coords are ordered slowest-to-fastest (e.g. (k, i, j)); the affine matrix
        # columns are fastest-to-slowest (x, y, z), so index it in reverse.
        val = data[n, dim, ndim]  # translation column
        for r in range(ndim):
            coord = starts[r] + coords[r] * steps[r]
            val += data[n, dim, ndim - 1 - r] * coord
        return val

    oshape = (data.shape[0], ndim, *target_shape)
    return te.compute(oshape, _compute, tag="affine_grid")


def _grid_sample_2d(
    data, grid, method="bilinear", layout="NCHW", padding_mode="zeros", align_corners=True
):
    """Applies bilinear/nearest/bicubic sampling to input feature map.

    Given :math:`data` and :math:`grid` assuming NCHW layout, then the output is computed by

    .. math::

        x_{src} = grid[batch, 0, y_{dst}, x_{dst}] \\
        y_{src} = grid[batch, 1, y_{dst}, x_{dst}] \\
        output[batch, channel, y_{dst}, x_{dst}] = G(data[batch, channel, y_{src}, x_{src})

    :math:`x_{dst}`, :math:`y_{dst}` enumerate all spatial locations in :math:`output`, and
    :math:`G()` denotes the interpolation method.

    The out-boundary points will be padded with zeros if padding_mode is "zeros", or
    border pixel value if padding_mode is "border", or
    inner pixel value if padding_mode is "reflection".

    The left-top corner (-1, -1) and right-bottom corner (1, 1) in grid will be map to
    (0, 0) and (h - 1, w - 1) of data if align_corners is "True", or
    (-0.5, -0.5) and (h - 0.5, w - 0.5) of data if align_corners is "False".

    The shape of the output will be (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]).

    The operator assumes that :math:`grid` has been normalized to [-1, 1].

    grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.

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

    grid : tvm.Tensor
        4-D with shape [batch, 2, out_height, out_width]

    method : str
        The interpolation method "nearest", "bilinear", "bicubic" are supported.

    layout : str
        The layout of input data and the output.

    padding_mode : str
        The padding mode for outside grid values, "zeros", "border", "reflection" are supported.

    align_corners: bool
        Geometrically, we consider the pixels of the input as squares rather than points.
        If set to "True", the extrema ("-1" and "1") are considered as referring
        to the center points of the input corner pixels. If set to "False", they
        are instead considered as referring to the corner points of the input corner
        pixels, making the sampling more resolution agnostic.

    Returns
    -------
    Output : tvm.Tensor
        4-D with shape [batch, in_channel, out_height, out_width]
    """

    assert method in ("bilinear", "nearest", "bicubic"), f"{method} is not supported"
    assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
    assert layout == "NCHW", f"{layout} is not supported"

    batch, in_channel, in_height, in_width = data.shape
    out_height, out_width = grid.shape[2:]

    def _get_pixel_value(n, c, h, w):
        return te.if_then_else(
            te.all(h >= 0, w >= 0, h < in_height, w < in_width),
            data[n, c, h, w],
            tirx.const(0.0, dtype=data.dtype),
        )

    def _unnormalize(h, w):
        if align_corners:
            y = (h + 1) * (in_height - 1) / 2
            x = (w + 1) * (in_width - 1) / 2
        else:
            y = -0.5 + (h + 1) * in_height / 2
            x = -0.5 + (w + 1) * in_width / 2
        return (y, x)

    def _clip_coordinates(x, size):
        return te.min(te.max(x, 0), size - 1)

    def _compute_source_index(n, h, w):
        y = grid[n, 1, h, w]
        x = grid[n, 0, h, w]
        y, x = _unnormalize(y, x)

        if padding_mode == "reflection":
            y = _reflect_coordinates(y, in_height)
            x = _reflect_coordinates(x, in_width)
            y = _clip_coordinates(y, in_height)
            x = _clip_coordinates(x, in_width)
        elif padding_mode == "border":
            y = _clip_coordinates(y, in_height)
            x = _clip_coordinates(x, in_width)

        return (y, x)

    def _reflect_coordinates(x, size):
        def __refelection(x, size, corner_start):
            def __reflect(index, size, corner_start):
                index_align_corner = te.abs(corner_start - index)
                size_times = te.truncdiv(index_align_corner.astype("int32"), size).astype("int32")
                t = tirx.Mod(size_times, 2)
                extra = index_align_corner - size_times * size
                return tirx.if_then_else(
                    tirx.EQ(t, 0), extra + corner_start, size - extra + corner_start
                )

            return tirx.if_then_else(
                tirx.all(x >= corner_start, x <= size + corner_start),
                x,
                __reflect(x, size, corner_start),
            )

        if align_corners:
            new_x = __refelection(x, size - 1, 0)
        else:
            new_x = __refelection(x, size, -0.5)
        return new_x

    def _bilinear_sample(n, c, h, w):
        y, x = _compute_source_index(n, h, w)
        y0 = te.floor(y).astype("int32")
        x0 = te.floor(x).astype("int32")
        y1 = y0 + tirx.const(1, "int32")
        x1 = x0 + tirx.const(1, "int32")

        return (
            _get_pixel_value(n, c, y0, x0) * (1.0 - (y - y0)) * (1.0 - (x - x0))
            + _get_pixel_value(n, c, y0, x1) * (1.0 - (y - y0)) * (x - x0)
            + _get_pixel_value(n, c, y1, x0) * (y - y0) * (1.0 - (x - x0))
            + _get_pixel_value(n, c, y1, x1) * (y - y0) * (x - x0)
        )

    def _nearest_sample(n, c, h, w):
        y, x = _compute_source_index(n, h, w)
        y_new = te.nearbyint(y).astype("int32")
        x_new = te.nearbyint(x).astype("int32")

        return _get_pixel_value(n, c, y_new, x_new)

    def _bicubic_sample(n, c, h, w):
        A = -0.75  # -0.75 is used in pytorch, it maybe different in other frameworks

        def cubic_weight_1(fraction):
            return ((A + 2) * fraction - (A + 3)) * fraction * fraction + 1

        def cubic_weight_2(fraction):
            return ((A * fraction - 5 * A) * fraction + 8 * A) * fraction - 4 * A

        def cubic_interp_1d(pixel_0, pixel_1, pixel_2, pixel_3, fraction):
            weights = [0] * 4
            weights[0] = cubic_weight_2(fraction + 1)
            weights[1] = cubic_weight_1(fraction)
            weights[2] = cubic_weight_1(1 - fraction)
            weights[3] = cubic_weight_2(2 - fraction)
            return (
                pixel_0 * weights[0]
                + pixel_1 * weights[1]
                + pixel_2 * weights[2]
                + pixel_3 * weights[3]
            )

        y = grid[n, 1, h, w]
        x = grid[n, 0, h, w]
        y, x = _unnormalize(y, x)
        y_floor = te.floor(y).astype("int32")
        x_floor = te.floor(x).astype("int32")
        y_fraction = y - y_floor
        x_fraction = x - x_floor

        coefficients = [0] * 4

        for i in range(4):
            y_ = y_floor - 1 + i
            x_0 = x_floor - 1
            x_1 = x_floor + 0
            x_2 = x_floor + 1
            x_3 = x_floor + 2

            if padding_mode == "border":
                y_ = _clip_coordinates(y_, in_height).astype("int32")
                x_0 = _clip_coordinates(x_0, in_width).astype("int32")
                x_1 = _clip_coordinates(x_1, in_width).astype("int32")
                x_2 = _clip_coordinates(x_2, in_width).astype("int32")
                x_3 = _clip_coordinates(x_3, in_width).astype("int32")

            elif padding_mode == "reflection":
                y_ = _reflect_coordinates(y_, in_height)
                x_0 = _reflect_coordinates(x_0, in_width)
                x_1 = _reflect_coordinates(x_1, in_width)
                x_2 = _reflect_coordinates(x_2, in_width)
                x_3 = _reflect_coordinates(x_3, in_width)

                y_ = _clip_coordinates(y_, in_height).astype("int32")
                x_0 = _clip_coordinates(x_0, in_width).astype("int32")
                x_1 = _clip_coordinates(x_1, in_width).astype("int32")
                x_2 = _clip_coordinates(x_2, in_width).astype("int32")
                x_3 = _clip_coordinates(x_3, in_width).astype("int32")

            coefficients[i] = cubic_interp_1d(
                _get_pixel_value(n, c, y_, x_0),
                _get_pixel_value(n, c, y_, x_1),
                _get_pixel_value(n, c, y_, x_2),
                _get_pixel_value(n, c, y_, x_3),
                x_fraction,
            )

        return cubic_interp_1d(
            coefficients[0], coefficients[1], coefficients[2], coefficients[3], y_fraction
        )

    if method == "bilinear":
        interpolation = _bilinear_sample
    elif method == "nearest":
        interpolation = _nearest_sample
    else:  # method == "bicubic"
        interpolation = _bicubic_sample

    return te.compute((batch, in_channel, out_height, out_width), interpolation, tag="grid_sample")


def _grid_sample_3d(
    data, grid, method="bilinear", layout="NCDHW", padding_mode="zeros", align_corners=True
):
    """Applies bilinear/nearest sampling to input feature map.

    Given :math:`data` and :math:`grid` assuming NCDHW layout, then the output is computed by

    .. math::

        x_{src} = grid[batch, 0, z_{dst}, y_{dst}, x_{dst}] \\
        y_{src} = grid[batch, 1, z_{dst}, y_{dst}, x_{dst}] \\
        z_{src} = grid[batch, 2, z_{dst}, y_{dst}, x_{dst}] \\
        output[batch, channel, z_{src}, y_{dst}, x_{dst}]
        = G(data[batch, channel, z_{src}, y_{src}, x_{src})

    :math:`x_{dst}`, :math:`y_{dst}`, :math:`z_{dst}` enumerate all spatial locations
    in :math:`output`, and :math:`G()` denotes the interpolation method.

    The out-boundary points will be padded with zeros if padding_mode is "zeros", or
    border pixel value if padding_mode is "border", or
    inner pixel value if padding_mode is "reflection".

    The left-top corner (-1, -1, -1) and right-bottom corner (1, 1, 1) in grid will be map to
    (0, 0, 0) and (d - 1, h - 1, w - 1) of data if align_corners is "True", or
    (-0.5, -0.5, -0.5) and (d - 0.5, h - 0.5, w - 0.5) of data if align_corners is "False".

    The shape of the output will be
    (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3], grid.shape[4]).

    The operator assumes that :math:`grid` has been normalized to [-1, 1].

    grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.

    Parameters
    ----------
    data : tvm.Tensor
        5-D with shape [batch, in_channel, in_depth, in_height, in_width]

    grid : tvm.Tensor
        5-D with shape [batch, 3, out_depth, out_height, out_width]

    method : str
        The interpolation method "nearest", "bilinear"("trilinear") are supported.

    layout : str
        The layout of input data and the output.

    padding_mode : str
        The padding mode for outside grid values, "zeros", "border", "reflection" are supported.

    align_corners: bool
        Geometrically, we consider the pixels of the input as squares rather than points.
        If set to "True", the extrema ("-1" and "1") are considered as referring
        to the center points of the input corner pixels. If set to "False", they
        are instead considered as referring to the corner points of the input corner
        pixels, making the sampling more resolution agnostic.

    Returns
    -------
    Output : tvm.Tensor
        5-D with shape [batch, in_channel, out_depth, out_height, out_width]
    """

    assert method in ("bilinear", "nearest"), f"{method} is not supported"
    assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
    assert layout == "NCDHW", f"{layout} is not supported"

    batch, in_channel, in_depth, in_height, in_width = data.shape
    out_depth, out_height, out_width = grid.shape[2:]

    def _get_pixel_value(n, c, d, h, w):
        return te.if_then_else(
            te.all(d >= 0, h >= 0, w >= 0, d < in_depth, h < in_height, w < in_width),
            data[n, c, d, h, w],
            tirx.const(0.0, dtype=data.dtype),
        )

    def _compute_source_index(n, d, h, w):
        z = grid[n, 2, d, h, w]
        y = grid[n, 1, d, h, w]
        x = grid[n, 0, d, h, w]

        if align_corners:
            z = (z + 1) * (in_depth - 1) / 2
            y = (y + 1) * (in_height - 1) / 2
            x = (x + 1) * (in_width - 1) / 2
        else:
            z = -0.5 + (z + 1) * in_depth / 2
            y = -0.5 + (y + 1) * in_height / 2
            x = -0.5 + (x + 1) * in_width / 2

        if padding_mode == "reflection":
            z = _reflect_coordinates(z, in_depth)
            y = _reflect_coordinates(y, in_height)
            x = _reflect_coordinates(x, in_width)
            z = _clip_coordinates(z, in_depth)
            y = _clip_coordinates(y, in_height)
            x = _clip_coordinates(x, in_width)
        elif padding_mode == "border":
            z = _clip_coordinates(z, in_depth)
            y = _clip_coordinates(y, in_height)
            x = _clip_coordinates(x, in_width)

        return (z, y, x)

    def _clip_coordinates(x, size):
        return te.min(te.max(x, 0), size - 1)

    def _reflect_coordinates(x, size):
        def __refelection(x, size, corner_start):
            def __reflect(index, size, corner_start):
                index_align_corner = te.abs(corner_start - index)
                size_times = te.truncdiv(index_align_corner.astype("int32"), size).astype("int32")
                t = tirx.Mod(size_times, 2)
                extra = index_align_corner - size_times * size
                return tirx.if_then_else(
                    tirx.EQ(t, 0), extra + corner_start, size - extra + corner_start
                )

            return tirx.if_then_else(
                tirx.all(x >= corner_start, x <= size + corner_start),
                x,
                __reflect(x, size, corner_start),
            )

        if align_corners:
            return __refelection(x, size - 1, 0)
        return __refelection(x, size, -0.5)

    def _trilinear_sample(n, c, d, h, w):
        z, y, x = _compute_source_index(n, d, h, w)
        z0 = te.floor(z).astype("int32")
        y0 = te.floor(y).astype("int32")
        x0 = te.floor(x).astype("int32")
        z1 = z0 + tirx.const(1, "int32")
        y1 = y0 + tirx.const(1, "int32")
        x1 = x0 + tirx.const(1, "int32")

        return (
            _get_pixel_value(n, c, z0, y0, x0) * (1 - (x - x0)) * (1 - (y - y0)) * (1 - (z - z0))
            + _get_pixel_value(n, c, z0, y0, x1) * (x - x0) * (1 - (y - y0)) * (1 - (z - z0))
            + _get_pixel_value(n, c, z1, y1, x0) * (1 - (x - x0)) * (y - y0) * (z - z0)
            + _get_pixel_value(n, c, z1, y1, x1) * (x - x0) * (y - y0) * (z - z0)
            + _get_pixel_value(n, c, z0, y1, x0) * (1 - (x - x0)) * (y - y0) * (1 - (z - z0))
            + _get_pixel_value(n, c, z1, y0, x1) * (x - x0) * (1 - (y - y0)) * (z - z0)
            + _get_pixel_value(n, c, z1, y0, x0) * (1 - (x - x0)) * (1 - (y - y0)) * (z - z0)
            + _get_pixel_value(n, c, z0, y1, x1) * (x - x0) * (y - y0) * (1 - (z - z0))
        )

    def _nearest_sample(n, c, d, h, w):
        z, y, x = _compute_source_index(n, d, h, w)
        z_new = te.nearbyint(z).astype("int32")
        y_new = te.nearbyint(y).astype("int32")
        x_new = te.nearbyint(x).astype("int32")
        return _get_pixel_value(n, c, z_new, y_new, x_new)

    if method == "bilinear":
        interpolation = _trilinear_sample
    else:  # method == "nearest"
        interpolation = _nearest_sample

    return te.compute(
        (batch, in_channel, out_depth, out_height, out_width), interpolation, tag="grid_sample"
    )


def grid_sample(
    data, grid, method="bilinear", layout="NCHW", padding_mode="zeros", align_corners=True
):
    """Applies grid sampling to input feature map.

    Given :math:`data` and :math:`grid`, then for 4-D the output is computed by

    .. math::

        x_{src} = grid[batch, 0, y_{dst}, x_{dst}] \\
        y_{src} = grid[batch, 1, y_{dst}, x_{dst}] \\
        output[batch, channel, y_{dst}, x_{dst}] = G(data[batch, channel, y_{src}, x_{src}])

    :math:`x_{dst}`, :math:`y_{dst}` enumerate all spatial locations in :math:`output`, and
    :math:`G()` denotes the interpolation function.

    The out-boundary points will be padded with zeros if padding_mode is "zeros", or
    border pixel value if padding_mode is "border", or
    inner pixel value if padding_mode is "reflection".

    The left-top corner (-1, -1) and right-bottom corner (1, 1) in grid will be map to
    (0, 0) and (h - 1, w - 1) of data if align_corners is "True", or
    (-0.5, -0.5) and (h - 0.5, w - 0.5) of data if align_corners is "False".

    The shape of the output will be
    4-D (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]), or
    5-D (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3], grid.shape[4]).

    The operator assumes that :math:`grid` has been normalized to [-1, 1].

    grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.

    Parameters
    ----------
    data : tvm.Tensor
        4-D with shape [batch, in_channel, in_height, in_width], or
        5-D with shape [batch, in_channel, in_depth, in_height, in_width]

    grid : tvm.Tensor
        4-D with shape [batch, 2, out_height, out_width], or
        5-D with shape [batch, 3, out_depth, out_height, out_width]

    method : str
        The interpolation method, 4-D "nearest", "bilinear", "bicubic" and
        5-D "nearest", "bilinear"("trilinear") are supported.

    layout : str
        The layout of input data and the output.

    padding_mode : str
        The padding mode for outside grid values, "zeros", "border", "reflection" are supported.

    align_corners: bool
        Geometrically, we consider the pixels of the input as squares rather than points.
        If set to "True", the extrema ("-1" and "1") are considered as referring
        to the center points of the input corner pixels. If set to "False", they
        are instead considered as referring to the corner points of the input corner
        pixels, making the sampling more resolution agnostic.

    Returns
    -------
    Output : tvm.Tensor
        4-D with shape [batch, in_channel, out_height, out_width], or
        5-D with shape [batch, in_channel, out_depth, out_height, out_width]
    """

    if len(layout) == 4:
        compute = _grid_sample_2d
    elif len(layout) == 5:
        compute = _grid_sample_3d
    else:
        msg = f"layout {layout} is not supported"
        raise ValueError(msg)

    return compute(data, grid, method, layout, padding_mode, align_corners)