# 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.
# ruff: noqa: RUF005
"""TVM operator pixel shuffle compute."""

import tvm


def pixel_shuffle(data, upscale_factor, name="PixelShuffle"):
    """PixelShuffle operator that rearranges elements in a tensor of shape
    [..., C * r * r, H, W] to [..., C, H * r, W * r].

    Parameters
    ----------
    data : tvm.te.Tensor
        N-D input tensor with at least 3 dimensions. Channel must be at index -3.

    upscale_factor : int
        The upscale factor (r).

    name : str
        Name of the output tensor.

    Returns
    -------
    output : tvm.te.Tensor
        Pixel shuffled tensor with shape [..., C, H*r, W*r]
    """
    assert isinstance(upscale_factor, int) and upscale_factor > 0
    ndim = len(data.shape)
    assert ndim >= 3, "Input must be at least 3D"

    upscale_factor_const = tvm.tirx.const(upscale_factor, "int32")
    c_in, h_in, w_in = data.shape[-3], data.shape[-2], data.shape[-1]

    c_out = tvm.tirx.floordiv(c_in, upscale_factor_const * upscale_factor_const)
    h_out = h_in * upscale_factor_const
    w_out = w_in * upscale_factor_const

    out_shape = list(data.shape[:-3]) + [c_out, h_out, w_out]

    def _compute(*indices):
        batch_indices = indices[:-3]
        c_out_idx, h_out_idx, w_out_idx = indices[-3], indices[-2], indices[-1]

        h_idx = tvm.tirx.floordiv(h_out_idx, upscale_factor_const)
        h_offset = h_out_idx % upscale_factor_const

        w_idx = tvm.tirx.floordiv(w_out_idx, upscale_factor_const)
        w_offset = w_out_idx % upscale_factor_const

        c_in_idx = (
            (c_out_idx * upscale_factor_const * upscale_factor_const)
            + (h_offset * upscale_factor_const)
            + w_offset
        )

        index_tuple = batch_indices + (c_in_idx, h_idx, w_idx)
        return data[index_tuple]

    return tvm.te.compute(out_shape, _compute, name=name)