# 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
"""ROI Align operator"""

import tvm
from tvm import te

from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc


def _sample_common(
    i,
    c,
    ph,
    pw,
    rois,
    pooled_size_h,
    pooled_size_w,
    spatial_scale,
    sample_ratio,
    aligned,
    dtype,
    avg_mode,
    bilinear_func,
):
    roi = rois[i]
    batch_index = roi[0].astype("int32")
    roi_start_w = roi[1] * spatial_scale
    roi_start_h = roi[2] * spatial_scale
    roi_end_w = roi[3] * spatial_scale
    roi_end_h = roi[4] * spatial_scale

    if aligned:
        roi_h = roi_end_h - roi_start_h
        roi_w = roi_end_w - roi_start_w
    else:
        roi_h = te.max(roi_end_h - roi_start_h, tvm.tirx.const(1.0, dtype))
        roi_w = te.max(roi_end_w - roi_start_w, tvm.tirx.const(1.0, dtype))

    pooled_size_h_const = tvm.tirx.const(pooled_size_h, dtype)
    pooled_size_w_const = tvm.tirx.const(pooled_size_w, dtype)
    bin_h = roi_h / pooled_size_h_const
    bin_w = roi_w / pooled_size_w_const

    if sample_ratio > 0:
        roi_bin_grid_h = tvm.tirx.const(sample_ratio, "int32")
        roi_bin_grid_w = tvm.tirx.const(sample_ratio, "int32")
    else:
        roi_bin_grid_h = te.ceil(roi_h / pooled_size_h_const).astype("int32")
        roi_bin_grid_w = te.ceil(roi_w / pooled_size_w_const).astype("int32")

    count = roi_bin_grid_h * roi_bin_grid_w
    rh = te.reduce_axis((0, roi_bin_grid_h), name="rh")
    rw = te.reduce_axis((0, roi_bin_grid_w), name="rw")
    roi_start_h = roi_start_h + tvm.tirx.Cast(dtype, ph) * bin_h
    roi_start_w = roi_start_w + tvm.tirx.Cast(dtype, pw) * bin_w

    def sample_value(rh_idx, rw_idx):
        return bilinear_func(
            batch_index,
            c,
            roi_start_h
            + (tvm.tirx.Cast(dtype, rh_idx) + tvm.tirx.const(0.5, dtype))
            * bin_h
            / tvm.tirx.Cast(dtype, roi_bin_grid_h),
            roi_start_w
            + (tvm.tirx.Cast(dtype, rw_idx) + tvm.tirx.const(0.5, dtype))
            * bin_w
            / tvm.tirx.Cast(dtype, roi_bin_grid_w),
        )

    if avg_mode:
        return te.sum(
            sample_value(rh, rw) / tvm.tirx.Cast(dtype, count),
            axis=[rh, rw],
        )
    return te.max(sample_value(rh, rw), axis=[rh, rw])


def roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1, aligned=False):
    """ROI align operator in NCHW layout."""
    avg_mode = mode in (b"avg", "avg", 0)
    max_mode = mode in (b"max", "max", 1)
    assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode."

    _, channel, height, width = data.shape
    num_roi, _ = rois.shape
    dtype = rois.dtype

    if isinstance(pooled_size, int):
        pooled_size_h = pooled_size_w = pooled_size
    else:
        pooled_size_h, pooled_size_w = pooled_size

    height_f = tvm.tirx.Cast(dtype, height)
    width_f = tvm.tirx.Cast(dtype, width)
    zero = tvm.tirx.const(0.0, data.dtype)

    def _bilinear(n, c, y, x):
        outside = tvm.tirx.any(y < -1.0, x < -1.0, y > height_f, x > width_f)
        y = te.min(te.max(y, 0.0), tvm.tirx.Cast(dtype, height - 1))
        x = te.min(te.max(x, 0.0), tvm.tirx.Cast(dtype, width - 1))
        val = bilinear_sample_nchw(data, (n, c, y, x), height - 1, width - 1)
        return tvm.tirx.if_then_else(outside, zero, val)

    return te.compute(
        (num_roi, channel, pooled_size_h, pooled_size_w),
        lambda i, c, ph, pw: _sample_common(
            i,
            c,
            ph,
            pw,
            rois,
            pooled_size_h,
            pooled_size_w,
            spatial_scale,
            sample_ratio,
            aligned,
            dtype,
            avg_mode,
            _bilinear,
        ),
        tag="pool,roi_align_nchw",
    )


def roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1, aligned=False):
    """ROI align operator in NHWC layout."""
    avg_mode = mode in (b"avg", "avg", 0)
    max_mode = mode in (b"max", "max", 1)
    assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode."

    _, height, width, channel = data.shape
    num_roi, _ = rois.shape
    dtype = rois.dtype

    if isinstance(pooled_size, int):
        pooled_size_h = pooled_size_w = pooled_size
    else:
        pooled_size_h, pooled_size_w = pooled_size

    height_f = tvm.tirx.Cast(dtype, height)
    width_f = tvm.tirx.Cast(dtype, width)
    zero = tvm.tirx.const(0.0, data.dtype)

    def _bilinear(n, c, y, x):
        outside = tvm.tirx.any(y < -1.0, x < -1.0, y > height_f, x > width_f)
        y = te.min(te.max(y, 0.0), tvm.tirx.Cast(dtype, height - 1))
        x = te.min(te.max(x, 0.0), tvm.tirx.Cast(dtype, width - 1))
        val = bilinear_sample_nhwc(data, (n, y, x, c), height - 1, width - 1)
        return tvm.tirx.if_then_else(outside, zero, val)

    return te.compute(
        (num_roi, pooled_size_h, pooled_size_w, channel),
        lambda i, ph, pw, c: _sample_common(
            i,
            c,
            ph,
            pw,
            rois,
            pooled_size_h,
            pooled_size_w,
            spatial_scale,
            sample_ratio,
            aligned,
            dtype,
            avg_mode,
            _bilinear,
        ),
        tag="pool,roi_align_nhwc",
    )


def roi_align(
    data,
    rois,
    pooled_size,
    spatial_scale,
    mode="avg",
    sample_ratio=-1,
    aligned=False,
    layout="NCHW",
):
    """ROI align operator."""
    if layout == "NCHW":
        return roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio, aligned)
    if layout == "NHWC":
        return roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio, aligned)
    raise ValueError(f"Unsupported layout for roi_align: {layout}")
