# 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.
"""Image operators."""

from typing import cast

from tvm import DataType
from tvm.ir import is_prim_expr

from ...expr import Expr, ShapeExpr
from . import _ffi_api

PrimExprLike = int | Expr
SizeLike = PrimExprLike | tuple[PrimExprLike, ...]


def resize2d(
    data: Expr,
    size: SizeLike,
    roi: float | tuple[float] | None = None,
    layout: str = "NCHW",
    method: str = "linear",
    coordinate_transformation_mode: str = "half_pixel",
    rounding_method: str = "round",
    cubic_alpha: float = -0.75,
    cubic_exclude: int = 0,
    extrapolation_value: float = 0.0,
    out_dtype: str | DataType | None = None,
) -> Expr:
    """Image resize2d operator.

    This operator takes data as input and does 2D scaling to the given scale factor.
    In the default case, where the data_layout is `NCHW`
    with data of shape (n, c, h, w)
    out will have a shape (n, c, size[0], size[1])

    method indicates the algorithm to be used while calculating the out value
    and method can be one of ("linear", "nearest_neighbor", "cubic")

    Parameters
    ----------
    data : relax.Expr
        The input data to the operator.

    size: SizeLike
        The out size to which the image will be resized.
        If specified as a list, it is required to have length either 1 or 2.
        If specified as an Expr, it is required to have ndim 2.

    roi: Optional[Union[float, Tuple[float]]]
        The region of interest for cropping the input image. Expected to be of
        size 4, and format [start_h, start_w, end_h, end_w].
        Only used if coordinate_transformation_mode is tf_crop_and_resize.

    layout : str
        Layout of the input.

    method : str
        Scale method to used [nearest_neighbor, linear, cubic].

    coordinate_transformation_mode : str
        Describes how to transform the coordinate in the resized tensor
        to the coordinate in the original tensor. Definitions can be found
        in topi/image/resize.py.
        [half_pixel, align_corners, asymmetric, pytorch_half_pixel,
        tf_half_pixel_for_nn, and tf_crop_and_resize].

    rounding_method: str
        indicates how to find the "nearest" pixel in nearest_neighbor method
        [round, floor, ceil]

    cubic_alpha: float
        Spline Coefficient for bicubic interpolation

    cubic_exclude: int
        Flag to exclude exterior of the image during bicubic interpolation

    extrapolation_value: float
        Fill value to use when roi is outside of the image

    out_dtype : Optional[str | DataType]
        The dtype of the output tensor.
        It it is not specified, the output will have the same dtype as input if not specified.

    Returns
    -------
    result: relax.Expr
        The resized result.
    """
    if roi is None:
        roi = (0.0, 0.0, 0.0, 0.0)  # type: ignore
    elif isinstance(roi, float):
        roi = (roi, roi, roi, roi)  # type: ignore
    elif isinstance(roi, tuple | list):
        roi = tuple(val if isinstance(val, float) else float(val) for val in roi)
    else:
        raise NotImplementedError(f"Unsupported roi type {type(roi)}")

    if isinstance(size, int) or is_prim_expr(size):
        size = (size, size)
    if isinstance(size, tuple | list):
        if len(size) == 1:
            size = ShapeExpr([size[0], size[0]])
        else:
            size = ShapeExpr(size)

    return _ffi_api.resize2d(  # type: ignore
        data,
        size,
        roi,
        layout,
        method,
        coordinate_transformation_mode,
        rounding_method,
        cubic_alpha,
        cubic_exclude,
        extrapolation_value,
        out_dtype,
    )


def resize3d(
    data: Expr,
    size: SizeLike,
    roi: float | tuple[float] | None = None,
    layout: str = "NCDHW",
    method: str = "linear",
    coordinate_transformation_mode: str = "half_pixel",
    rounding_method: str = "",
    cubic_alpha: float = -0.75,
    cubic_exclude: int = 0,
    extrapolation_value: float = 0.0,
    out_dtype: str | DataType | None = None,
) -> Expr:
    """Image resize3d operator.

    This operator takes data as input and does 3D scaling to the given output size.
    In the default case, where data layout is `NCDHW`
    with data of shape (n, c, d, h, w),
    the output has shape (n, c, size[0], size[1], size[2]).
    """
    if roi is None:
        roi = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)  # type: ignore
    elif isinstance(roi, float):
        roi = (roi, roi, roi, roi, roi, roi)  # type: ignore
    elif isinstance(roi, tuple | list):
        roi = tuple(val if isinstance(val, float) else float(val) for val in roi)
    else:
        raise NotImplementedError(f"Unsupported roi type {type(roi)}")

    if isinstance(size, int) or is_prim_expr(size):
        size = (size, size, size)
    if isinstance(size, tuple | list):
        if len(size) == 1:
            size = ShapeExpr([size[0], size[0], size[0]])
        else:
            size = ShapeExpr(size)

    return _ffi_api.resize3d(  # type: ignore
        data,
        size,
        roi,
        layout,
        method,
        coordinate_transformation_mode,
        rounding_method,
        cubic_alpha,
        cubic_exclude,
        extrapolation_value,
        out_dtype,
    )


def grid_sample(
    data: Expr,
    grid: Expr,
    method: str = "bilinear",
    layout: str = "NCHW",
    padding_mode: str = "zeros",
    align_corners: bool = False,
) -> Expr:
    """Applies grid sampling to input feature map.

    Given data and grid, the output is computed by sampling from data using
    the grid coordinates.

    Parameters
    ----------
    data : relax.Expr
        The input data tensor with shape [N, C, H, W] for NCHW layout.

    grid : relax.Expr
        The grid tensor with shape [N, H_out, W_out, 2]. The values are normalized
        to [-1, 1], where (-1, -1) is the top-left corner and (1, 1) is the bottom-right.

    method : str
        Interpolation method. Can be 'nearest', 'bilinear', or 'bicubic'.

    layout : str
        Layout of the input data. Default is 'NCHW'.

    padding_mode : str
        Padding mode for outside grid values. Can be 'zeros', 'border', or 'reflection'.

    align_corners : bool
        If True, the corner pixels of the input and output tensors are aligned.

    Returns
    -------
    result : relax.Expr
        The sampled output tensor with shape [N, C, H_out, W_out].
    """
    return _ffi_api.grid_sample(  # type: ignore
        data,
        grid,
        method,
        layout,
        padding_mode,
        align_corners,
    )


def affine_grid(
    data: Expr,
    size: SizeLike,
    align_corners: bool = True,
) -> Expr:
    """Generate a 2D or 3D sampling grid using an affine transformation matrix.

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

    Parameters
    ----------
    data : relax.Expr
        The input affine matrix tensor with shape [batch, 2, 3] for 2D or
        [batch, 3, 4] for 3D.

    size : SizeLike
        The target output spatial shape, (H, W) for 2D or (D, H, W) for 3D. If a
        single integer or PrimExpr is provided, it is interpreted as a square 2D
        output shape (size, size).

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

    Returns
    -------
    result : relax.Expr
        The output grid tensor with shape [batch, 2, H, W] for 2D or
        [batch, 3, D, H, W] for 3D.
    """
    if isinstance(size, int) or is_prim_expr(size):
        size = (size, size)
    if isinstance(size, tuple | list):
        size = ShapeExpr(size)

    return cast(Expr, _ffi_api.affine_grid(data, size, align_corners))
