# 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=import-outside-toplevel, redefined-builtin, unused-argument
"""Set operators."""

import numpy as np  # type: ignore

import tvm

from ..expr import Expr, prim_value
from . import _ffi_api


def unique(
    x: Expr,
    sorted: bool | Expr = True,
    return_index: bool | Expr = False,
    return_inverse: bool | Expr = False,
    return_counts: bool | Expr = False,
    axis: int | Expr | None = None,
) -> Expr:
    """Find the unique elements in a given tensor.
    In addition, it optionally returns
    - the indices of the input tensor that give the unique values;
    - the indices of the unique tensor that reconstruct the input tensor;
    - the number of times each unique value comes up in the input tensor.

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

    sorted : Union[bool, Expr]
        Whether to sort the unique elements in ascending order before
        returning as output.

    return_index : Union[bool, Expr]
        Whether to return an additional tensor with indices for where elements in
        the unique tensor come from the original input.

    return_inverse : Union[bool, Expr]
        Whether to return an additional tensor with indices for where elements in
        the original input ended up in the returned unique list.

    return_counts : Union[bool, Expr]
        Whether to return an additional tensor with counts of each unique elements.

    axis : Optional
        The dimension to apply unique.
        If not specified, the unique values of the flattened input are returned.

    Returns
    -------
    ret : relax.Expr
        The created relax call with
    """

    if isinstance(sorted, bool):
        sorted = prim_value(sorted)
    if isinstance(return_index, bool):
        return_index = prim_value(return_index)
    if isinstance(return_inverse, bool):
        return_inverse = prim_value(return_inverse)
    if isinstance(return_counts, bool):
        return_counts = prim_value(return_counts)
    if axis is not None and isinstance(axis, int):
        axis = prim_value(axis)
    return _ffi_api.unique(  # type: ignore
        x, sorted, return_index, return_inverse, return_counts, axis
    )


@tvm.register_global_func("relax.run.unique")
def numpy_unique(
    x: tvm.runtime.tensor,
    sorted: int,
    return_index: int,
    return_inverse: int,
    return_counts: int,
    axis: int | None = None,
) -> tvm.runtime.tensor:
    """Returns the unique elements of the input tensor.

    Uses numpy.unique to compute unique elements.
    """
    import builtins

    x_numpy = x.numpy()

    # Call numpy.unique with all the requested return flags
    result = np.unique(
        x_numpy,
        return_index=bool(return_index),
        return_inverse=bool(return_inverse),
        return_counts=bool(return_counts),
        axis=axis,
    )

    # If no optional outputs requested, result is just the unique values
    if not bool(return_index) and not bool(return_inverse) and not bool(return_counts):
        unique_values = result
        if not sorted:
            indices = np.unique(x_numpy, return_index=True, axis=axis)[1]
            unique_values = np.take(x_numpy, builtins.sorted(indices), axis=axis)
        return tvm.runtime.tensor(unique_values)

    # Otherwise, numpy returns a tuple
    unique_values = result[0]
    output_list = []
    result_idx = 1

    # Handle sorting for unique values
    if not sorted and bool(return_index):
        # Get the indices from numpy result
        indices = result[result_idx]
        result_idx += 1
        # Sort indices to get original order
        sort_order = np.argsort(indices)
        unique_values = np.take(unique_values, sort_order, axis=axis)
        indices = np.sort(indices)
        output_list.append(tvm.runtime.tensor(unique_values))
        output_list.append(tvm.runtime.tensor(indices))
    elif not sorted:
        # Need to get indices to reorder
        _, indices = np.unique(x_numpy, return_index=True, axis=axis)
        sort_order = np.argsort(indices)
        unique_values = np.take(unique_values, sort_order, axis=axis)
        output_list.append(tvm.runtime.tensor(unique_values))
        if bool(return_index):
            indices_from_result = result[result_idx]
            result_idx += 1
            output_list.append(tvm.runtime.tensor(np.sort(indices_from_result)))
    else:
        # Sorted case
        output_list.append(tvm.runtime.tensor(unique_values))
        if bool(return_index):
            output_list.append(tvm.runtime.tensor(result[result_idx]))
            result_idx += 1

    if bool(return_inverse):
        inverse_indices = result[result_idx]
        if not sorted:
            # Need to remap inverse indices to match reordered unique values
            _, orig_indices = np.unique(x_numpy, return_index=True, axis=axis)
            sort_order = np.argsort(orig_indices)
            inverse_mapping = np.empty_like(sort_order)
            inverse_mapping[sort_order] = np.arange(len(sort_order))
            inverse_indices = inverse_mapping[inverse_indices]
        # ONNX spec: inverse_indices is always 1D
        # When axis is None, it has length X.size (flattened)
        # When axis is specified, it has length X.shape[axis]
        # numpy.unique already returns 1D inverse_indices, so no reshaping needed
        output_list.append(tvm.runtime.tensor(inverse_indices))
        result_idx += 1

    if bool(return_counts):
        counts = result[result_idx]
        if not sorted:
            # Reorder counts to match reordered unique values
            _, orig_indices = np.unique(x_numpy, return_index=True, axis=axis)
            sort_order = np.argsort(orig_indices)
            counts = counts[sort_order]
        output_list.append(tvm.runtime.tensor(counts))

    return tuple(output_list)


def nonzero(x: Expr) -> Expr:
    """Find the indices of elements of a tensor that are non-zero.

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

    Returns
    -------
    result : relax.Expr
        A 2-D tensor containing indices of non-zero elements.

    Note
    ----
    This function is equivalent to `onnx.nonzero`.

    Examples
    --------

    .. code-block:: python

        x = [[0, 1],
             [2, 0]]
        nonzero(x) = [[0, 1],
                      [1, 0]]

    """
    return _ffi_api.nonzero(x)  # type: ignore


@tvm.register_global_func("relax.run.nonzero")
def numpy_nonzero(x: tvm.runtime.tensor) -> tvm.runtime.tensor:
    np_result = np.atleast_1d(x.numpy()).nonzero()
    return tvm.runtime.tensor(np.stack(np_result, axis=0))
