# 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,consider-using-enumerate,redefined-outer-name
"""Injective transformation operators"""

from math import pi

import numpy as np

import tvm
from tvm import te, topi

from . import cpp, tag
from .utils import const_vector, make_idx, within_index


def expand_dims(a, axis, num_newaxis=1):
    """Expand the shape of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be expanded.

    num_newaxis: int, optional
        Number of newaxis to be inserted on axis

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.expand_dims(a, axis, num_newaxis)


def expand_like(a, shape_like, axis):
    """Expand an input array with the shape of second array.
    This operation can always be composed of unsqueezing and
    expanding dims on those unsqueezed axes.

    Examples
    --------
    .. code-block::

        input = [ 12.  19.  27.]
        input.shape = (3,)

        new_shape_array = [[[1,2],[2,3],[1,3]],
                        [[1,4],[4,3],[5,2]],
                        [[7,1],[7,2],[7,3]]]
        new_shape_array.shape = (3, 3, 2)

        expand_like(input, [1,2], new_shape_array) =
                        [[[12,12],[12,12],[12,12]],
                        [[19,19],[19,19],[19,19]],
                        [[27,27],[27,27],[27,27]]]

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be expanded.
    shape_like : tvm.te.Tensor
        The tensor to with target shape.
    axis: list of int
        axis to be expanded on

    Returns
    -------
    ret : tvm.te.Tensor
    """
    odim = len(axis) + len(a.shape)
    if odim != len(shape_like.shape):
        if len(a.shape) == 1 and len(axis) == len(shape_like.shape):
            # A special case: `a` is a scalar represented as a 1-dim tensor
            return te.compute(shape_like.shape, lambda *idxs: a(0))
        raise ValueError(
            f"shape inconsistent when expand_like ({len(axis)}, "
            f"{len(a.shape)}, {len(shape_like.shape)})"
        )

    real_axis = topi.reduction._get_real_axis(len(shape_like.shape), axis)
    real_axis = sorted(real_axis)

    def _compute(*idxs):
        indices = []
        axis_index = 0
        for i in range(0, len(idxs)):
            if i not in real_axis:
                dim = tvm.tirx.if_then_else(a.shape[len(indices)] != 1, idxs[i], 0)
                indices.append(dim)
                axis_index += 1
        return a(*indices)

    return te.compute(shape_like.shape, _compute)


def transpose(a, axes=None):
    """Permute the dimensions of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be expanded.

    axes: tuple of ints, optional
        By default, reverse the dimensions.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.transpose(a, axes)


def flip(a, axis=0):
    """Flip/reverse elements of an array in a particular axis.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be expanded.

    axis : int, optional
        The axis along which the tensors will be reveresed.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.flip(a, axis)


def reverse_sequence(a, seq_lengths, seq_axis=1, batch_axis=0):
    """Reverse the tensor for variable length slices.
    Input is first sliced along batch axis and then elements are reversed along seq axis.

    Parameters
    ----------
    a : tvm.te.Tensor
       The tensor to be reversed.

    seq_lengths : tvm.te.Tensor
       A 1D Tensor with length a.dims[batch_axis]
       Must be one of the following types: int32, int64
       if seq_lengths[i] > a.dims[seq_axis], it is rounded to a.dims[seq_axis]
       if seq_lengths[i] < 1, it is rounded to 1

    seq_axis : int, optional
       The axis along which the elements will be reversed. Default is 1.

    batch_axis : int, optional
       The axis along which the tensor will be sliced. Default is 0.

    Returns
    -------
    ret : tvm.te.Tensor
       The computed result of same shape and type as of input.

    """
    return cpp.reverse_sequence(a, seq_lengths, seq_axis, batch_axis)


def strided_slice(a, begin, end, strides=None, axes=None, slice_mode="end", assume_inbound=True):
    """Slice of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be sliced.

    begin : list of int
        The indices to begin with in the slicing.

    end : list of int
        Indices indicating end of the slice.

    strides : list of int, optional
        Specifies the stride values, it can be negative
        in that case, the input tensor will be reversed
        in that particular axis.

    axes : list of int, optional
        Axes along which slicing is applied. When it is specified, begin, end
        strides, and axes need to a list of integers of the same length.

    slice_mode : str, optional
        The slice mode [end, size].
        end - The ending indices for the slice [default].
        size - The input strides will be ignored, input end in this mode indicates
        the sizeof a slice starting at the location specified by begin. If end[i]
        is -1, all remaining elements in that dimension are included in the slice.

    assume_inbound: bool, optional
        A flag to indicate if all indices are assumed to be inbound

    Returns
    -------
    ret : tvm.te.Tensor
    """
    if (
        isinstance(begin, tvm.te.Tensor)
        or isinstance(end, tvm.te.Tensor)
        or isinstance(strides, tvm.te.Tensor)
    ):
        assert axes is None, "axes argument is not supported by dynamic strided slice yet."
        if not isinstance(begin, tvm.te.Tensor):
            begin = const_vector(begin)
        if not isinstance(end, tvm.te.Tensor):
            end = const_vector(end)
        if strides is None:
            strides = [1] * begin.shape[0].value
        if not isinstance(strides, tvm.te.Tensor):
            strides = const_vector(strides)
        return cpp.dynamic_strided_slice(a, begin, end, strides)
    if strides is None:
        strides = []
    if axes is None:
        axes = []
    # axes is a list of host integers on the C++ side (Array<int64_t>); unwrap any
    # IntImm entries that callers may pass through (e.g. relax legalize pipeline).
    axes = [int(v) if isinstance(v, tvm.tirx.IntImm) else v for v in axes]
    return cpp.strided_slice(a, begin, end, strides, axes, slice_mode, assume_inbound)


def dynamic_strided_slice(a, begin, end, strides, output_shape):
    """Slice of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be sliced.

    begin : tvm.te.Tensor
        The indices to begin with in the slicing.

    end : tvm.te.Tensor
        Indices indicating end of the slice.

    strides : tvm.te.Tensor
        Specifies the stride values, it can be negative
        in that case, the input tensor will be reversed
        in that particular axis.

    output_shape: list of Expr
        Specifies the output shape

    Returns
    -------
    ret : tvm.te.Tensor
    """
    if not isinstance(begin, tvm.te.Tensor):
        begin = const_vector(begin)
    if not isinstance(end, tvm.te.Tensor):
        end = const_vector(end)
    if not isinstance(strides, tvm.te.Tensor):
        strides = const_vector(strides)
    return cpp.relax_dynamic_strided_slice(a, begin, end, strides, output_shape)


@tvm.te.tag_scope(tag=tag.INJECTIVE + ",strided_set")
def strided_set(a, v, begin, end, strides=None):
    """Set slice of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be sliced.

    v : tvm.te.Tensor
        The values to set

    begin: tvm.te.Tensor
        The indices to begin with in the slicing.

    end: tvm.te.Tensor
        Indices indicating end of the slice.

    strides: tvm.te.Tensor, optional
        Specifies the stride values, it can be negative
        in that case, the input tensor will be reversed
        in that particular axis.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    n = len(a.shape)

    if len(begin.shape) != 1:
        raise ValueError("begin should be a vector")
    if not begin.dtype == "int32":
        raise TypeError("begin should be int32")
    if len(end.shape) != 1:
        raise ValueError("end should be a vector")
    if not end.dtype == "int32":
        raise TypeError("end should be int32")
    if strides is not None:
        if len(strides.shape) != 1:
            raise ValueError("strides should be a vector")
        if not strides.dtype == "int32":
            raise TypeError("strides should be int32")

    def _max(a, b):
        return tvm.tirx.Select(a > b, a, b)

    if strides is None:
        strides = [tvm.tirx.const(1, "int32")] * n
    else:
        strides = [
            tvm.tirx.if_then_else(strides.shape[0] > i, strides[i], tvm.tirx.const(1, "int32"))
            for i in range(n)
        ]

    begin = [
        tvm.tirx.if_then_else(
            begin.shape[0] > i,
            begin[i],
            tvm.tirx.Select(strides[i] > 0, tvm.tirx.const(0, "int32"), a.shape[i]),
        )
        for i in range(n)
    ]
    end = [
        tvm.tirx.if_then_else(
            end.shape[0] > i,
            end[i],
            tvm.tirx.Select(strides[i] > 0, a.shape[i] + 1, -(a.shape[i] + 1)),
        )
        for i in range(n)
    ]

    # Convert negative indexes
    for i in range(n):
        begin[i] = tvm.tirx.if_then_else(begin[i] < 0, begin[i] + a.shape[i], begin[i])
        end[i] = tvm.tirx.if_then_else(end[i] < 0, end[i] + a.shape[i], end[i])

    def _select(*indices):
        from_val = []
        index_tuple = []
        for i in range(n):
            from_val.append(within_index(begin[i], end[i], strides[i], indices[i]))
            index_tuple.append(make_idx(begin[i], end[i], strides[i], a.shape[i], indices[i]))
        return tvm.tirx.if_then_else(tvm.tirx.all(*from_val), v(*index_tuple), a(*indices))

    return te.compute(a.shape, _select, name="strided_set")


def reshape(a, newshape):
    """Reshape the array

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be reshaped
    newshape : tuple of ints
        The new shape

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.reshape(a, newshape)


def squeeze(a, axis=None):
    """Remove single-dimensional entries from the shape of an array.

    Parameters
    ----------
    a : tvm.te.Tensor

    axis : None or int or tuple of ints, optional
        Selects a subset of the single-dimensional entries in the shape.
        If an axis is selected with shape entry greater than one, an error is raised.

    Returns
    -------
    squeezed : tvm.te.Tensor
    """
    return cpp.squeeze(a, axis)


def concatenate(a_tuple, axis=0):
    """Join a sequence of arrays along an existing axis.

    Parameters
    ----------
    a_tuple : tuple of tvm.te.Tensor
        The arrays to concatenate

    axis : int, optional
        The axis along which the arrays will be joined. Default is 0.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.concatenate(a_tuple, axis)


def stack(tensors, axis=0):
    """Join a sequence of tensors along a new axis.

    Parameters
    ----------
    tensors : tuple or list of tvm.te.Tensor
        The tensors to be stacked. All tensors must have the same shape.

    axis : int, optional
        The axis in the resulting tensor along which the input tensors will be stacked.
        Negative values wrap around. Default is 0.

    Returns
    -------
    ret : tvm.te.Tensor
        The stacked tensor with an additional dimension compared to the input tensors.

    """
    return cpp.stack(tensors, axis)


def split(ary, indices_or_sections, axis=0):
    """Split an array into multiple sub-arrays.

    Parameters
    ----------
    ary : tvm.te.Tensor

    indices_or_sections : int or 1-D array

    axis : int

    Returns
    -------
    ret : tuple of tvm.te.Tensor
    """
    return cpp.split(ary, indices_or_sections, axis)


def take(a, indices, axis=None, batch_dims=0, mode="fast"):
    """Take elements from an array along an axis.

    Parameters
    ----------
    a : tvm.te.Tensor
        The source array.

    indices : tvm.te.Tensor
        The indices of the values to extract.

    axis : int, optional
        The axis over which to select values. By default,
        the flattened input array is used.

    batch_dims : int
        The number of batch dimensions. By default is 0.

    mode : str, optional
        Specifies how out-of-bounds indices will behave.
        - fast (default): extra indices lead to seg fault (user must make sure indices are in-bound)
        - nan: produce NaNs for out-of-bounds indices
        - wrap: wrap around the indices
        - clip: clip to the range

    Returns
    -------
    ret : tvm.te.Tensor
    """
    if axis is None:
        return cpp.take(a, indices, int(batch_dims), mode)
    return cpp.take(a, indices, int(batch_dims), int(axis), mode)


def gather(data, axis, indices):
    """Gather values along given axis from given indices.

    E.g. for a 3D tensor, output is computed as:

    .. code-block:: python

        out[i][j][k] = data[indices[i][j][k]][j][k]  # if axis == 0
        out[i][j][k] = data[i][indices[i][j][k]][k]  # if axis == 1
        out[i][j][k] = data[i][j][indices[i][j][k]]  # if axis == 2

    ``indices`` must have same shape as ``data``, except at dimension ``axis``
    which must just be not null. Output will have same shape as ``indices``.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input data to the operator.

    axis: int
        The axis along which to index.

    indices : tvm.te.Tensor
        The indices of the values to extract.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.gather(data, axis, indices)


def gather_nd(a, indices, batch_dims=0):
    """Gather elements from a n-dimension array..

    Parameters
    ----------
    a : tvm.te.Tensor
        The source array.

    indices : tvm.te.Tensor
        The indices of the values to extract.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.gather_nd(a, indices, batch_dims)


def matmul(a, b, transp_a=False, transp_b=False):
    """
    Creates an operation that calculates a matrix multiplication (row-major notation):
    A(i, k) * B(k, j)
    if trans_a == trans_b, the usual transposed combinations, otherwise

    Parameters
    ----------
    a : The matrix A
    b : The matrix B
    trans_a : Is A's layout transposed?
    trans_b : Is B's layout transposed?

    Returns
    -------
    A Tensor whose op member is the matmul operation
    """
    return cpp.matmul(a, b, transp_a, transp_b)


def tensordot(a, b, axes):
    """A generalization of matrix multiplication to tensor.

    Parameters
    ----------
    a : The tensor A
    b : The tensor B
    axes : The number of dimensions to reduce over

    Returns
    -------
    A Tensor computing the result
    """
    if isinstance(axes, int):
        return cpp.tensordot(a, b, axes)
    if isinstance(axes[0], int):
        return cpp.tensordot(a, b, (axes[0],), (axes[1],))
    return cpp.tensordot(a, b, axes[0], axes[1])


def arange(start, stop=None, step=1, dtype="float32"):
    """Creates a tensor with evenly spaced values within a given interval.

    Parameters
    ----------
    start : tvm.Expr, optional
        Start of interval. The interval includes this value. The default start
        value is 0.

    stop : tvm.Expr
        Stop of interval. The interval does not include this value.

    step : tvm.Expr, optional
        Spacing between values. The default step size is 1.

    dtype : str, optional
        The target data type.

    Returns
    -------
    result : tvm.te.Tensor
        The resulting tensor.
    """
    if stop is None:
        stop = start
        start = 0
    return cpp.arange(start, stop, step, dtype)


def meshgrid(a_tuple, indexing):
    """Create coordinate matrices from coordinate vectors.

    Parameters
    ----------
    a_tuple : tuple of tvm.te.Tensor
        The coordinate vectors or scalars.

    indexing : str
        Indexing mode, either "ij" or "xy".

    Returns
    -------
    result : tuple of tvm.te.Tensor
        The resulting grids for each axis.
    """
    return cpp.meshgrid(a_tuple, indexing)


def repeat(a, repeats, axis):
    """Repeats elements of an array.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be repeated.

    repeats: int, required
        Number of repetitions for each element

    axis: int, optional
        The axis along which to repeat values

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.repeat(a, repeats, axis)


def tile(a, reps):
    """Repeats the whole array multiple times.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be tiled.

    reps: tuple of ints, required
        The number of times for repeating the tensor

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.tile(a, reps)


def dyn_tile(a, new_shape, rdim):
    """Repeats the whole array multiple times with dynamic output shape.

    Parameters
    ----------
    a : tvm.te.Tensor
        The tensor to be tiled.

    new_shape : tuple of Expr
        The output shape after tiling.

    rdim : int
        The rank of the repeats input.

    Returns
    -------
    ret : tvm.te.Tensor
    """

    return cpp.dyn_tile(a, new_shape, rdim)


def layout_transform(array, src_layout, dst_layout, schedule_rule="None"):
    """Transform the layout according to src_layout and dst_layout

    Parameters
    ----------
    array : tvm.te.Tensor
        The source array.

    src_layout : str
        the source layout.

    dst_layout : str
        the destination layout.

    schedule_rule : str
        the schedule rule to apply if any
    """
    return cpp.layout_transform(array, src_layout, dst_layout, schedule_rule)


def shape(array, dtype="int32"):
    """Get the shape of input array

    Parameters
    ----------
    array : tvm.te.Tensor
        The source tensor.

    dtype : str, optional
        The target data type.

    Returns
    -------
    result : tvm.te.Tensor
        The resulting tensor.
    """
    return cpp.shape(array, dtype)


def sequence_mask(data, valid_length, mask_value=0, axis=0):
    """Sets all elements outside the expected length of the sequence to a constant value.

    This function takes an n-dimensional input array of the form [MAX_LENGTH, batch_size, ...] or
    [batch_size, MAX_LENGTH, ...] and returns an array of the same shape.

    `axis` means the axis of the length dimension and can only be 0 or 1. If `axis` is 0,
    the data must have shape [MAX_LENGTH, batch_size, ...]. Otherwise (axis=1), the data must have
    shape [batch_size, MAX_LENGTH, ...].

    `valid_length` gives the length of each sequence. `valid_length` should be
    a 1D int array with positive ints and has dimension [batch_size,].

    Parameters
    ----------
    data : tvm.te.Tensor
        N-D with shape [MAX_LENGTH, batch_size, ...] or [batch_size, MAX_LENGTH, ...]
        depending on the value of `axis`.

    valid_length : tvm.te.Tensor
        1-D with shape [batch_size,]

    mask_value : float, optional
        The masking value, default 0

    axis : int, optional
        axis of the length dimension, must be 0 or 1, default 0

    Returns
    -------
    output : tvm.te.Tensor
        N-D with shape [MAX_LENGTH, batch_size, ...] or [batch_size, MAX_LENGTH, ...]
        depending on the value of `axis`.
    """

    assert len(data.shape) >= 2, f"only support data.ndim >= 2, received data.shape = {data.shape}"
    assert axis in (0, 1), f"only support axis = 0, 1, received axis = {axis}"
    return cpp.sequence_mask(data, valid_length, mask_value, axis)


def tensor_size(array, dtype="int32"):
    """Get the number of elements of input array

    Parameters
    ----------
    array : tvm.te.Tensor
        The source tensor.

    dtype : str, optional
        The target data type.

    Returns
    -------
    result : tvm.te.Tensor
        The resulting tensor.
    """
    return cpp.tensor_size(array, dtype)


def where(condition, x, y):
    """Get the elements, either from x or y, depending on the condition.

    Parameters
    ----------
    condition : tvm.te.Tensor
        The condition array.

    x : tvm.te.Tensor
        First array to be selected.

    y : tvm.te.Tensor
        Second array to be selected.

    Returns
    -------
    result : tvm.te.Tensor
        A Tensor selected from x or y depending on condition.
    """
    return cpp.where(condition, x, y)


def one_hot(indices, on_value, off_value, depth, axis, dtype):
    """
    Returns a one-hot tensor where the locations repsented by indices take value on_value,
    other locations take value off_value.
    Final dimension is <indices outer dimensions> x depth x <indices inner dimensions>.

    Parameters
    ----------
    indices : tvm.te.Tensor
        Locations to set to on_value.

    on_value : tvm.te.Tensor
        Value to fill at indices.

    off_value : tvm.te.Tensor
        Value to fill at all other positions besides indices.

    depth : int
        Depth of the one-hot dimension.

    axis : int
        Axis to fill.

    dtype : str
        Data type of the output tensor.

    Returns
    -------
    ret : tvm.te.Tensor
        The one-hot tensor.

    Examples
    --------
    .. code-block:: python

        indices = [0, 1, 2]

        topi.one_hot(indices, 3) =
            [[1, 0, 0],
             [0, 1, 0],
             [0, 0, 1]]
    """
    return cpp.one_hot(indices, on_value, off_value, depth, axis, dtype)


def unravel_index(indices, shape):
    """Convert a flat index or array of flat indices into a tuple of coordinate arrays.

    Example::
    -   unravel_index([22, 41, 37], [7, 6]) = [[3, 6, 6], [4, 5, 1]]

    Parameters
    ----------
    indices : tvm.te.Tensor
        An integer array containing indices.

    shape : tvm.te.Tensor
        The shape of the array.

    Returns
    -------
    result : tvm.te.Tensor
        The tuple of coordinate arrays.
    """

    return cpp.unravel_index(indices, shape)


def sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value=0):
    """Converts a sparse representation into a dense tensor.

    Example::
    -   sparse_to_dense([[0, 0], [1, 1]], [2, 2], [3, 3], 0) = [[3, 0], [0, 3]]

    Parameters
    ----------
    sparse_indices : tvm.te.Tensor
        A 0-D, 1-D, or 2-D tensor of integers containing location of sparse values.

    output_shape : A list of integers
        Shape of the dense output tensor.

    sparse_values : tvm.te.Tensor
        A 0-D or 1-D tensor containing the sparse values for the sparse indices.

    default_value : tvm.te.Tensor
        A 0-D tensor containing the default value for the remaining locations.
        Defaults to 0.

    Returns
    -------
    result : tvm.te.Tensor
        Dense tensor of shape output_shape. Has the same type as sparse_values.
    """

    return cpp.sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value)


def matrix_set_diag(data, diagonal, k=0, align="RIGHT_LEFT"):
    """
    Returns a tensor with the diagonals of input tensor replaced with the provided diagonal values.

    Parameters
    ----------
    data : tvm.te.Tensor
        Input Tensor.

    diagonal : tvm.te.Tensor
        Values to be filled in the diagonal.

    k : int or tuple of int, optional
        Diagonal Offset(s). The diagonal or range of diagonals to set. (0 by default)
        Positive value means superdiagonal, 0 refers to the main diagonal, and
        negative value means subdiagonals. k can be a single integer (for a single diagonal)
        or a pair of integers specifying the low and high ends of a matrix band.
        k[0] must not be larger than k[1].

    align : string, optional
        Some diagonals are shorter than max_diag_len and need to be padded.
        align is a string specifying how superdiagonals and subdiagonals should be aligned,
        respectively. There are four possible alignments: "RIGHT_LEFT" (default), "LEFT_RIGHT",
        "LEFT_LEFT", and "RIGHT_RIGHT". "RIGHT_LEFT" aligns superdiagonals to the right
        (left-pads the row) and subdiagonals to the left (right-pads the row). It is the packing
        format LAPACK uses. cuSPARSE uses "LEFT_RIGHT", which is the opposite alignment.

    Returns
    -------
    result : tvm.te.Tensor
        New tensor with given diagonal values.

    Examples
    --------
    .. code-block:: python

        data = [[[7, 7, 7, 7],
                 [7, 7, 7, 7],
                 [7, 7, 7, 7]],
                [[7, 7, 7, 7],
                 [7, 7, 7, 7],
                 [7, 7, 7, 7]]]

        diagonal = [[1, 2, 3],
                    [4, 5, 6]]

        topi.matrix_set_diag(input, diagonal) =
            [[[1, 7, 7, 7],
              [7, 2, 7, 7],
              [7, 7, 3, 7]],
             [[4, 7, 7, 7],
              [7, 5, 7, 7],
              [7, 7, 6, 7]]]
    """
    if isinstance(k, tuple | list):
        k_one = k[0]
        if len(k) >= 2:
            k_two = k[1]
        else:
            k_two = k[0]
    else:
        k_one = k
        k_two = k

    super_diag_right_align = align[:5] == "RIGHT"
    sub_diag_right_align = align[-5:] == "RIGHT"

    return cpp.matrix_set_diag(
        data, diagonal, k_one, k_two, super_diag_right_align, sub_diag_right_align
    )


def adv_index(data, indices):
    """Numpy style indexing with tensors.

    Parameters
    ----------
    data : tvm.te.Tensor
        Input data.

    indices : A list of tvm.te.Tensor
        Tensor index.

    Returns
    -------
    result : tvm.te.Tensor
        Output tensor
    """
    return cpp.adv_index(data, indices)


def sliding_window(data, axis, window_shape, strides):
    """Slide a window over the data tensor.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input data to the operator.

    axis : int
        What axis the window begins sliding over. Window will be slid over
        this axis and all following axes. The axis value determines the window
        shape (and thus, the number of strides): window shape and strides must
        both be of length `data.ndim-axis`.

    window_shape : List[int]
        The window shape to form over the input. Window shape must be of length
        `data.ndim-axis`.

    strides : List[int]
        How to stride the window along each dimension. Strides must be of length
        `data.ndim-axis`.

    Returns
    -------
    result : tvm.te.Tensor
        The resulting tensor.
    """
    return cpp.sliding_window(data, axis, window_shape, strides)


def trilu(data, k, upper):
    """
    Given a 2-D matrix or batches of 2-D matrices, returns the
    upper or lower triangular part of the tensor.

    Parameters
    ----------
    data: tvm.te.Tensor
        The tensor that trilu will be applied to. Must be either
        a 2D matrix or a tensor of batches of 2D matrices.

    k: tvm.te.Tensor
        The number of diagonals above or below the main diagonal
        to exclude or include.

    upper: bool
        If True, only upper triangular values of input are kept,
        if False, the lower triangular values are kept.


    Returns
    -------
    ret : tvm.te.Tensor
        The new tensor with appropriate diagonals set to zero.

    Examples
    --------
    .. code-block:: python

        x = [[0, 1, 2],
             [3, 4, 5],
             [6, 7, 8]]

        topi.trilu(x, True, 0) =
            [[0, 1, 2],
             [0, 4, 5],
             [0, 0, 8]]
    """
    # Make sure datatype is consistent.
    if k.ty != tvm.ir.PrimType("int32"):
        k = tvm.tirx.Cast("int32", k)

    # Check either above or below diagonal depending on upper.
    check_op = tvm.tirx.GE
    if upper:
        check_op = tvm.tirx.LE

    def _apply_trilu(*indices):
        row_index = indices[-2]
        col_index = indices[-1]
        # promote row & col indices
        if row_index.ty != col_index.ty:
            target_type = (col_index + row_index).ty
            if row_index.ty != target_type:
                row_index = tvm.tirx.Cast(target_type, row_index)
            else:
                col_index = tvm.tirx.Cast(target_type, col_index)
        other_indices = indices[:-2]
        check_position = check_op(row_index, col_index - k)
        value = data(*other_indices, row_index, col_index)
        return tvm.tirx.Select(check_position, value, tvm.tirx.const(0, data.dtype))

    return te.compute(data.shape, _apply_trilu, name="trilu", tag=topi.tag.ELEMWISE)


def index_tensor(data, indices):
    """Advanced-tensor indexing (NumPy/PyTorch-style).

    Given k index tensors ``indices = (I0, I1, …, Ik-1)`` this
    operator selects elements from ``data`` as if one had written
    ``data[I0, I1, …, Ik-1]`` in NumPy/PyTorch:

    * All index tensors must have an integer dtype.
    * Their shapes are broadcast together to a common shape ``B`` in
      the usual NumPy way.
    * The result shape is ``B + data.shape[k:]`` (i.e. the broadcast
      shape followed by the remaining axes of ``data`` that are *not*
      indexed).
    *  ``k`` must not exceed ``data.ndim``; otherwise a compile-time
       error is raised.

    Parameters
    ----------
    data : tvm.te.Tensor
        The tensor to be indexed.

    indices : Sequence[tvm.te.Tensor]
        A Python ``list`` / ``tuple`` of **k** index tensors,
        or a `tvm.te.Tensor` tuple expression. Each tensor must have an
        integer dtype.

    Returns
    -------
    result : tvm.te.Tensor
        The tensor obtained after advanced indexing.  Its dtype equals
        ``data.dtype``

    Examples
    --------
    .. code-block:: python

        x     = te.placeholder((3, 3),  name="x")        # shape (3,3)
        row   = te.placeholder((2,),    name="row", dtype="int32")
        col   = te.placeholder((2,),    name="col", dtype="int32")

        # Equivalent to x[row, col] in NumPy / PyTorch
        y = topi.index_tensor(x, [row, col])             # shape (2,)

        # Broadcasting example:
        row = te.placeholder((2, 1), name="row", dtype="int32")
        col = te.placeholder((1, 3), name="col", dtype="int32")
        z = topi.index_tensor(x, [row, col])             # shape (2, 3)
    """
    return topi.adv_index(data, indices)


def hamming_window(window_size, periodic, alpha, beta, dtype):
    """Hamming window function.

    Parameters
    ----------
    window_size: tvm.Expr
        The size of returned window.

    periodic: tvm.Expr
        If True, returns a window to be used as periodic function.
        If False, return a symmetric window.

    alpha: tvm.Expr
        The co-efficient alpha.

    beta: tvm.Expr
        The co-efficient beta.

    Returns
    -------
    ret : tvm.te.Tensor
        The result tensor.
    """
    if window_size == 1:
        return topi.const_vector(np.array([1], dtype=dtype))

    periodic = topi.cast(periodic, "bool")

    if periodic:
        window_size += 1

    index = topi.arange(0, window_size, dtype=dtype)
    angular_freq = 2 * pi * index / (window_size - 1)
    cos_values = topi.cos(angular_freq)
    window = topi.cast(alpha - beta * cos_values, dtype=dtype)

    if periodic:
        return topi.strided_slice(window, [0], [window.shape[0] - 1])

    return window
