# 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=redefined-builtin,consider-using-enumerate,no-member
"""Reduce operators"""

from . import cpp


def _get_real_axis(ndim, axis):
    if axis is None:
        real_axis = list(range(ndim))
    else:
        if isinstance(axis, int):
            axis = [axis]
        else:
            assert isinstance(axis, list | tuple)
        real_axis = []
        for ele in axis:
            if ele < 0:
                ele += ndim
            if ele >= ndim:
                raise ValueError(
                    f"{ele} exceeds the maximum dimension {ndim}. Received axis={axis}"
                )
            real_axis.append(ele)
        real_axis.sort()
        real_axis = list(set(real_axis))  # Remove the duplicates
    return real_axis


def sum(data, axis=None, keepdims=False):
    """Sum of array elements over a given axis or a list of axes

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm tensor

    axis : None or int or tuple of int
        Axis or axes along which a sum is performed.
        The default, axis=None, will sum all of the elements of the input array.
        If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

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


def all(data, axis=None, keepdims=False):
    """Logical AND of array elements over a given axis or a list of axes

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm boolean tensor

    axis : None or int or tuple of int
        Axis or axes along which a logical AND is performed.
        The default, axis=None, will perform logical AND over all elements of the input array.
        If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

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


def any(data, axis=None, keepdims=False):
    """Logical OR of array elements over a given axis or a list of axes

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm boolean tensor

    axis : None or int or tuple of int
        Axis or axes along which a logical OR is performed.
        The default, axis=None, will perform logical OR over all elements of the input array.
        If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

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


def max(data, axis=None, keepdims=False):
    """Maximum of array elements over a given axis or a list of axes

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm tensor

    axis : None or int or tuple of int
        Axis or axes along which the max operation is performed.
        The default, axis=None, will find the max element from all of the elements of the input
        array. If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

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


def min(data, axis=None, keepdims=False):
    """Minimum of array elements over a given axis or a list of axes

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm tensor

    axis : None or int or tuple of int
        Axis or axes along which a minimum operation is performed.
        The default, axis=None, will find the minimum element from all of the elements of the
        input array. If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

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


def argmax(data, axis=None, keepdims=False, select_last_index=False):
    """Returns the indices of the maximum values along an axis.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm tensor

    axis : None or int or tuple of int
        Axis or axes along which a argmax operation is performed.
        The default, axis=None, will find the indices of the maximum element of the elements of
        the input array. If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

    select_last_index: bool
        Whether to select the last index if the maximum element appears multiple times, else
        select the first index.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.argmax(data, axis, keepdims, select_last_index)


def argmin(data, axis=None, keepdims=False, select_last_index=False):
    """Returns the indices of the minimum values along an axis.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm tensor

    axis : None or int or tuple of int
        Axis or axes along which a argmin operation is performed.
        The default, axis=None, will find the indices of minimum element all of the elements of
        the input array. If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

    select_last_index: bool
        Whether to select the last index if the minimum element appears multiple times, else
        select the first index.

    Returns
    -------
    ret : tvm.te.Tensor
    """
    return cpp.argmin(data, axis, keepdims, select_last_index)


def prod(data, axis=None, keepdims=False):
    """Product of array elements over a given axis or a list of axes

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tvm tensor

    axis : None or int or tuple of int
        Axis or axes along which a prod operation is performed.
        The default, axis=None, will get the prod element over all of the elements of the
        input array. If axis is negative it counts from the last to the first axis.

    keepdims : bool
        If this is set to True, the axes which are reduced are left in the result as dimensions
        with size one.
        With this option, the result will broadcast correctly against the input array.

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


def collapse_sum(data, target_shape):
    """Return a summation of data to the given shape.

    collapse_sum is intended as the backward operator of topi broadcast operators in the automatic
    differentiation process.

    We expect that data is the result of broadcasting some tensor of target_shape in some
    broadcast operation. Thus target_shape and data.shape must follow broadcast rules.

    During computation, the axes of data.shape and target_shape are checked from right to left.
    For every axis, if it either:
    - exist in data but not in target_shape, or
    - is larger than 1 in data and equals to 1 in target_shape,
    data will be summed over this axis.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input tensor.

    shape : Tuple[int]
        The shape to collapse to.

    Returns
    -------
    ret : tvm.te.Tensor
        The result tensor after summation.
    """
    return cpp.collapse_sum(data, target_shape)
