# SPDX-License-Identifier: Apache-2.0
"""Reduction operators."""

from typing import TYPE_CHECKING

import onnx
import torch

from ..op_registry import register
from ..utils.attributes import get_attribute
from ..utils.op_helpers import get_attribute_or_input

if TYPE_CHECKING:
    from ..graph_builder import GraphBuilder


def _get_reduction_axes(
    node: onnx.NodeProto, builder: "GraphBuilder"
) -> list[int] | torch.fx.Node | None:
    """Get axes for reduction, handling both attribute and input formats.

    In opset < 13, axes is an attribute.
    In opset 13-17, axes can be an attribute or an optional input.
    In opset 18+, axes is an optional input only.
    """
    return get_attribute_or_input(
        builder,
        node,
        attr_name="axes",
        input_index=1,
        opset_version=builder.opset_version,
        attr_allowed_until=17,
        input_allowed_since=13,
        default=None,
    )


@register("ReduceSum")
def reduce_sum(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Sum reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)
    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)

    def _reduce_sum(t, axes, keepdims, noop_with_empty_axes):
        # Handle empty axes list
        if isinstance(axes, (list, tuple)) and len(axes) == 0:
            if noop_with_empty_axes:
                return t
            # Empty axes with noop=False means reduce all dimensions
            axes = None
        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
            if noop_with_empty_axes:
                return t
            axes = None
        if axes is None:
            result = torch.sum(t)
            if keepdims:
                # Reshape to have all dimensions as 1
                result = result.reshape([1] * t.ndim)
            return result
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.sum(t, dim=axes, keepdim=keepdims)

    return builder.call_function(
        _reduce_sum, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
    )


@register("ReduceMean")
def reduce_mean(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Mean reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)
    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)

    def _reduce_mean(t, axes, keepdims, noop_with_empty_axes):
        # Handle empty axes list
        if isinstance(axes, (list, tuple)) and len(axes) == 0:
            if noop_with_empty_axes:
                return t
            axes = None
        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
            if noop_with_empty_axes:
                return t
            axes = None
        if axes is None:
            result = torch.mean(t)
            if keepdims:
                result = result.reshape([1] * t.ndim)
            return result
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.mean(t, dim=axes, keepdim=keepdims)

    return builder.call_function(
        _reduce_mean, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
    )


@register("ReduceMax")
def reduce_max(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Max reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)
    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)

    def _reduce_max(t, axes, keepdims, noop_with_empty_axes):
        # Handle empty axes list
        if isinstance(axes, (list, tuple)) and len(axes) == 0:
            if noop_with_empty_axes:
                return t
            axes = None
        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
            if noop_with_empty_axes:
                return t
            axes = None

        if axes is None:
            # Reduce over all dimensions
            if t.numel() == 0:
                # Empty tensor: return -inf with proper shape
                if keepdims:
                    return torch.full(
                        [1] * t.ndim, float("-inf"), dtype=t.dtype, device=t.device
                    )
                return torch.tensor(float("-inf"), dtype=t.dtype, device=t.device)
            result = t.max()
            if keepdims:
                result = result.reshape([1] * t.ndim)
            return result

        if isinstance(axes, torch.Tensor):
            axes = axes.tolist()
        if isinstance(axes, list) and len(axes) == 1:
            axes = axes[0]
        if isinstance(axes, int):
            # Check for empty dimension
            if t.shape[axes] == 0:
                new_shape = list(t.shape)
                if keepdims:
                    new_shape[axes] = 1
                else:
                    new_shape.pop(axes)
                return torch.full(
                    new_shape, float("-inf"), dtype=t.dtype, device=t.device
                )
            return t.max(dim=axes, keepdim=keepdims).values
        # Multiple axes: reduce sequentially
        result = t
        for axis in sorted(axes, reverse=True):
            if result.shape[axis] == 0:
                new_shape = list(result.shape)
                if keepdims:
                    new_shape[axis] = 1
                else:
                    new_shape.pop(axis)
                result = torch.full(
                    new_shape, float("-inf"), dtype=result.dtype, device=result.device
                )
            else:
                result = result.max(dim=axis, keepdim=keepdims).values
        return result

    return builder.call_function(
        _reduce_max, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
    )


@register("ReduceMin")
def reduce_min(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Min reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)
    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)

    def _reduce_min(t, axes, keepdims, noop_with_empty_axes):
        # Handle empty axes list
        if isinstance(axes, (list, tuple)) and len(axes) == 0:
            if noop_with_empty_axes:
                return t
            axes = None
        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
            if noop_with_empty_axes:
                return t
            axes = None

        if axes is None:
            # Reduce over all dimensions
            if t.numel() == 0:
                # Empty tensor: return inf with proper shape
                if keepdims:
                    return torch.full(
                        [1] * t.ndim, float("inf"), dtype=t.dtype, device=t.device
                    )
                return torch.tensor(float("inf"), dtype=t.dtype, device=t.device)
            result = t.min()
            if keepdims:
                result = result.reshape([1] * t.ndim)
            return result

        if isinstance(axes, torch.Tensor):
            axes = axes.tolist()
        if isinstance(axes, list) and len(axes) == 1:
            axes = axes[0]
        if isinstance(axes, int):
            # Check for empty dimension
            if t.shape[axes] == 0:
                new_shape = list(t.shape)
                if keepdims:
                    new_shape[axes] = 1
                else:
                    new_shape.pop(axes)
                return torch.full(
                    new_shape, float("inf"), dtype=t.dtype, device=t.device
                )
            return t.min(dim=axes, keepdim=keepdims).values
        # Multiple axes: reduce sequentially
        result = t
        for axis in sorted(axes, reverse=True):
            if result.shape[axis] == 0:
                new_shape = list(result.shape)
                if keepdims:
                    new_shape[axis] = 1
                else:
                    new_shape.pop(axis)
                result = torch.full(
                    new_shape, float("inf"), dtype=result.dtype, device=result.device
                )
            else:
                result = result.min(dim=axis, keepdim=keepdims).values
        return result

    return builder.call_function(
        _reduce_min, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
    )


@register("ReduceProd")
def reduce_prod(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Product reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)
    noop_with_empty_axes = get_attribute(node, "noop_with_empty_axes", 0)

    def _reduce_prod(t, axes, keepdims, noop_with_empty_axes):
        # Handle empty axes list
        if isinstance(axes, (list, tuple)) and len(axes) == 0:
            if noop_with_empty_axes:
                return t
            axes = None
        if isinstance(axes, torch.Tensor) and axes.numel() == 0:
            if noop_with_empty_axes:
                return t
            axes = None

        if axes is None:
            result = torch.prod(t)
            if keepdims:
                result = result.reshape([1] * t.ndim)
            return result

        if isinstance(axes, torch.Tensor):
            axes = axes.tolist()
        if isinstance(axes, list) and len(axes) == 1:
            axes = axes[0]
        if isinstance(axes, int):
            return torch.prod(t, dim=axes, keepdim=keepdims)
        # Multiple axes: reduce sequentially
        result = t
        for axis in sorted(axes, reverse=True):
            result = torch.prod(result, dim=axis, keepdim=keepdims)
        return result

    return builder.call_function(
        _reduce_prod, args=(x, axes, bool(keepdims), bool(noop_with_empty_axes))
    )


@register("ReduceL1")
def reduce_l1(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """L1 norm reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)

    def _reduce_l1(t, axes, keepdims):
        abs_t = torch.abs(t)
        if axes is None:
            return torch.sum(abs_t)
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.sum(abs_t, dim=axes, keepdim=keepdims)

    return builder.call_function(_reduce_l1, args=(x, axes, bool(keepdims)))


@register("ReduceL2")
def reduce_l2(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """L2 norm reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)

    def _reduce_l2(t, axes, keepdims):
        if axes is None:
            return torch.norm(t)
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.norm(t, dim=axes, keepdim=keepdims)

    return builder.call_function(_reduce_l2, args=(x, axes, bool(keepdims)))


@register("ReduceLogSum")
def reduce_log_sum(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Log of sum reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)

    def _reduce_log_sum(t, axes, keepdims):
        if axes is None:
            return torch.log(torch.sum(t))
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.log(torch.sum(t, dim=axes, keepdim=keepdims))

    return builder.call_function(_reduce_log_sum, args=(x, axes, bool(keepdims)))


@register("ReduceLogSumExp")
def reduce_log_sum_exp(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """LogSumExp reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)

    def _reduce_log_sum_exp(t, axes, keepdims):
        if axes is None:
            return torch.logsumexp(t, dim=tuple(range(t.dim())))
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.logsumexp(t, dim=axes, keepdim=keepdims)

    return builder.call_function(_reduce_log_sum_exp, args=(x, axes, bool(keepdims)))


@register("ReduceSumSquare")
def reduce_sum_square(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Sum of squares reduction."""
    x = builder.get_value(node.input[0])
    axes = _get_reduction_axes(node, builder)
    keepdims = get_attribute(node, "keepdims", 1)

    def _reduce_sum_square(t, axes, keepdims):
        sq = torch.square(t)
        if axes is None:
            return torch.sum(sq)
        if isinstance(axes, torch.Tensor):
            axes = tuple(axes.tolist())
        elif isinstance(axes, (list, tuple)):
            axes = tuple(axes)
        return torch.sum(sq, dim=axes, keepdim=keepdims)

    return builder.call_function(_reduce_sum_square, args=(x, axes, bool(keepdims)))


# =============================================================================
# ArgMax/ArgMin operators
# =============================================================================


def _make_arg_extremum_handler(torch_fn):
    """Factory for ArgMax/ArgMin operator handlers."""

    def _arg_extremum(t, axis, keepdims, select_last_index):
        if select_last_index:
            flipped = torch.flip(t, [axis])
            idx = torch_fn(flipped, dim=axis, keepdim=keepdims)
            return t.size(axis) - 1 - idx
        return torch_fn(t, dim=axis, keepdim=keepdims)

    def handler(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
        x = builder.get_value(node.input[0])
        axis = get_attribute(node, "axis", 0)
        keepdims = get_attribute(node, "keepdims", 1)
        select_last_index = get_attribute(node, "select_last_index", 0)
        return builder.call_function(
            _arg_extremum, args=(x, axis, bool(keepdims), bool(select_last_index))
        )

    return handler


@register("ArgMax")
def argmax(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Index of maximum value."""
    return _make_arg_extremum_handler(torch.argmax)(builder, node)


@register("ArgMin")
def argmin(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Index of minimum value."""
    return _make_arg_extremum_handler(torch.argmin)(builder, node)


# =============================================================================
# Cumulative and TopK operators
# =============================================================================


@register("CumSum")
def cumsum(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Cumulative sum."""
    x = builder.get_value(node.input[0])
    axis = builder.get_value(node.input[1])

    exclusive = get_attribute(node, "exclusive", 0)
    reverse = get_attribute(node, "reverse", 0)

    def _cumsum(x, axis, exclusive, reverse):
        ax = axis.item() if isinstance(axis, torch.Tensor) else axis

        if reverse:
            x = torch.flip(x, [int(ax)])

        result = torch.cumsum(x, dim=int(ax))

        if exclusive:
            # Shift by one and pad with zero
            pad_shape = list(x.shape)
            pad_shape[int(ax)] = 1
            zero_pad = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
            result = torch.cat(
                [zero_pad, result.narrow(int(ax), 0, x.shape[int(ax)] - 1)], dim=int(ax)
            )

        if reverse:
            result = torch.flip(result, [int(ax)])

        return result

    return builder.call_function(_cumsum, args=(x, axis, exclusive, reverse))


@register("TopK")
def topk(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """Find top K values and indices."""
    x = builder.get_value(node.input[0])
    k = builder.get_value(node.input[1])

    axis = get_attribute(node, "axis", -1)
    largest = get_attribute(node, "largest", 1)
    sorted_ = get_attribute(node, "sorted", 1)

    def _topk(x, k, axis, largest, sorted_):
        k_val = k.item() if isinstance(k, torch.Tensor) else k
        k_val = int(k_val)

        # Handle unsupported dtypes (e.g., uint64) by converting to int64
        original_dtype = x.dtype
        needs_conversion = original_dtype == torch.uint64
        if needs_conversion:
            x = x.to(torch.int64)

        # ONNX TopK requires stable sorting: for equal values, the element
        # with lower index appears first. PyTorch's topk is not stable.
        # We achieve stability by using argsort on a composite key.
        # Create indices tensor for tie-breaking
        size = x.shape[axis]
        # Create indices [0, 1, 2, ..., size-1] along the specified axis
        indices_shape = [1] * x.ndim
        indices_shape[axis] = size
        idx = torch.arange(size, device=x.device, dtype=x.dtype).view(indices_shape)
        idx = idx.expand_as(x)

        # Scale values so that the index becomes the tiebreaker
        # For largest=True: negate values, sort ascending, lower index wins
        # For largest=False: use values directly, sort ascending, lower index wins
        if bool(largest):
            # Negate so that larger values become smaller (for ascending sort)
            # Add small offset based on index to break ties (lower index = smaller offset)
            sort_values = -x
        else:
            sort_values = x

        # Use argsort with stable=True for stable sorting
        sorted_indices = torch.argsort(sort_values, dim=axis, stable=True)

        # Take top k indices
        # Narrow to first k elements along axis
        top_k_indices = torch.narrow(sorted_indices, axis, 0, k_val)

        # Gather values using the indices
        values = torch.gather(x, axis, top_k_indices)

        # If sorted=False, the order is undefined, but we still use stable order
        # The indices should be the original indices
        indices = top_k_indices

        # Convert values back to original dtype if needed
        if needs_conversion:
            values = values.to(original_dtype)

        return values, indices

    result = builder.call_function(_topk, args=(x, k, axis, largest, sorted_))

    # Handle multiple outputs
    for i, output_name in enumerate(node.output):
        if output_name:
            idx_node = builder.call_function(lambda t, idx: t[idx], args=(result, i))
            builder.env[output_name] = idx_node

    return result
