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

from tvm import DataType
from tvm.ir import PrimType

from ..expr import Expr
from . import _ffi_api


def _raw_dtype(dtype):
    return dtype.dtype if isinstance(dtype, PrimType) else dtype


def max(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the max of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a max operation is performed.
        The default, axis=None, will compute the max of all elements in the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.max(x, axis, keepdims)  # type: ignore


def mean(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the mean of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a mean operation is performed.
        The default, axis=None, will compute the mean of all elements in the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.mean(x, axis, keepdims)  # type: ignore


def min(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the min of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a min operation is performed.
        The default, axis=None, will compute the min of all elements in the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.min(x, axis, keepdims)  # type: ignore


def prod(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the product of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a product is performed.
        The default, axis=None, will compute the product of all elements of the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.prod(x, axis, keepdims)  # type: ignore


def std(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the standard deviation of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a standard deviation is performed.
        The default, axis=None, will compute the std of all elements of the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.std(x, axis, keepdims)  # type: ignore


def sum(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the sum of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a sum is performed.
        The default, axis=None, will sum all of the elements of the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.sum(x, axis, keepdims)  # type: ignore


def cumprod(
    data: Expr,
    axis: int | None = None,
    dtype: str | DataType | None = None,
    exclusive: bool = False,
):
    """Numpy style cumprod op. Return the cumulative product of the elements along
    a given axis.

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

    axis : Optional[int]
        Axis along which the cumulative product is computed. The default (None) is to compute
        the cumprod over the flattened array.

    dtype : Optional[Union[str, DataType]]
        Type of the returned array and of the accumulator in which the elements are computed.
        If dtype is not specified, it defaults to the dtype of data.

    exclusive : bool
        If false (default), all elements are included in the product.  If
        true, the first element is excluded from the product.

    Returns
    -------
    result : relax.Expr
        The result has the same size as data, and the same shape as data if axis is not None.
        If axis is None, the result is a 1-d array.

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

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

        cumprod(a)  # if axis is not provided, cumprod is done over the flattened input.
        -> [ 1,  2,  6, 24, 120, 720]

        cumprod(a, dtype="float32")
        -> [  1.,  2.,  6., 24., 120., 720.]

        cumprod(a, axis=0)  # multiply over rows for each of the 3 columns
        -> [[1, 2, 3],
            [4, 10, 18]]

        cumprod(a, axis=1)
        -> [[ 1,  2,  6],
            [ 4,  20, 120]]

        a = [1, 1, 1, 0, 1, 1, 0]  # a is a boolean array
        cumprod(a, dtype=int32)  # dtype should be provided to get the expected results
        -> [1, 1, 1, 0, 0, 0, 0]
    """
    if exclusive is None:
        exclusive = False

    return _ffi_api.cumprod(data, axis, _raw_dtype(dtype), exclusive)  # type: ignore


def cumsum(
    data: Expr,
    axis: int | None = None,
    dtype: str | DataType | None = None,
    exclusive: bool = False,
):
    """Numpy style cumsum op. Return the cumulative inclusive sum of the elements along
    a given axis.

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

    axis : Optional[int]
        Axis along which the cumulative sum is computed. The default (None) is to compute
        the cumsum over the flattened array.

    dtype : Optional[Union[str, DataType]]
        Type of the returned array and of the accumulator in which the elements are summed.
        If dtype is not specified, it defaults to the dtype of data.

    exclusive : bool
        If false (default), all elements are included in the sum.  If
        true, the first element is excluded from the sum.

    Returns
    -------
    result : relax.Expr
        The result has the same size as data, and the same shape as data if axis is not None.
        If axis is None, the result is a 1-d array.

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

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

        cumsum(a)  # if axis is not provided, cumsum is done over the flattened input.
        -> [ 1,  3,  6, 10, 15, 21]

        cumsum(a, dtype="float32")
        -> [  1.,   3.,   6.,  10.,  15.,  21.]

        cumsum(a, axis=0)  # sum over rows for each of the 3 columns
        -> [[1, 2, 3],
            [5, 7, 9]]

        cumsum(a, axis=1)
        -> [[ 1,  3,  6],
            [ 4,  9, 15]]

        a = [1, 0, 1, 0, 1, 1, 0]  # a is a boolean array
        cumsum(a, dtype=int32)  # dtype should be provided to get the expected results
        -> [1, 1, 2, 2, 3, 4, 4]
    """
    if exclusive is None:
        exclusive = False

    return _ffi_api.cumsum(data, axis, _raw_dtype(dtype), exclusive)  # type: ignore


def variance(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the variance of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis or axes along which a variance operation is performed.
        The default, axis=None, will compute the variance of all elements in the input tensor.
        Negative indexing is supported.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.variance(x, axis, keepdims)  # type: ignore


def median(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False) -> Expr:
    """Computes the median of tensor elements over given axes.

    Parameters
    ----------
    x : relax.Expr
        The input data tensor

    axis : Optional[Union[int, List[int]]]
        Axis along which the median is computed. The default (None) is to compute
        the median of the entire flattened tensor.

    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 tensor.

    Returns
    -------
    result : relax.Expr
        The computed result.
    """
    if isinstance(axis, int):
        axis = [axis]
    return _ffi_api.median(x, axis, keepdims)  # type: ignore