# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

import awkward as ak
from awkward._connect.numpy import UNSUPPORTED
from awkward._dispatch import high_level_function
from awkward._layout import HighLevelContext
from awkward._namedaxis import (
    _get_named_axis,
    _keep_named_axis,
    _named_axis_to_positional_axis,
    _remove_named_axis,
)
from awkward._nplikes.numpy_like import NumpyMetadata
from awkward._regularize import regularize_axis

__all__ = ("max", "nanmax")

np = NumpyMetadata.instance()


@high_level_function()
def max(
    array,
    axis=None,
    *,
    keepdims=False,
    initial=None,
    mask_identity=True,
    highlevel=True,
    behavior=None,
    attrs=None,
):
    """
    Args:
        array: Array-like data (anything #ak.to_layout recognizes).
        axis (None or int or str): If None, combine all values from the array into
            a single scalar result; if an int, group by that axis: `0` is the
            outermost, `1` is the first level of nested lists, etc., and
            negative `axis` counts from the innermost: `-1` is the innermost,
            `-2` is the next level up, etc; if a str, it is interpreted as the
            name of the axis which maps to an int if named axes are present.
            Named axes are attached to an array using #ak.with_named_axis and
            removed with #ak.without_named_axis; also see the
            [Named axes user guide](../../user-guide/how-to-array-properties-named-axis.html).
        keepdims (bool): If False, this reducer decreases the number of
            dimensions by 1; if True, the reduced values are wrapped in a new
            length-1 dimension so that the result of this operation may be
            broadcasted with the original array.
        initial (None or number): The minimum value of an output element, as
            an alternative to the numeric type's natural identity (e.g. negative
            infinity for floating-point types, a minimum integer for integer types).
            If you use `initial`, you might also want `mask_identity=False`.
        mask_identity (bool): If True, reducing over empty lists results in
            None (an option type); otherwise, reducing over empty lists
            results in the operation's identity.
        highlevel (bool): If True, return an #ak.Array; otherwise, return
            a low-level #ak.contents.Content subclass.
        behavior (None or dict): Custom #ak.behavior for the output array, if
            high-level.
        attrs (None or dict): Custom attributes for the output array, if
            high-level.

    Returns the maximum value in each group of elements from `array` (many
    types supported, including all Awkward Arrays and Records). The identity
    of maximization is `-inf` if floating-point or the smallest integer value
    if applied to integers. This identity is usually masked: the maximum of
    an empty list is None, unless `mask_identity=False`.
    This operation is the same as NumPy's
    [amax](https://docs.scipy.org/doc/numpy/reference/generated/numpy.amax.html)
    if all lists at a given dimension have the same length and no None values,
    but it generalizes to cases where they do not.

    See #ak.sum for a more complete description of nested list and missing
    value (None) handling in reducers.

    See also #ak.nanmax.
    """
    # Dispatch
    yield (array,)

    # Implementation
    return _impl(
        array,
        axis,
        keepdims,
        initial,
        mask_identity,
        highlevel,
        behavior,
        attrs,
    )


@high_level_function()
def nanmax(
    array,
    axis=None,
    *,
    keepdims=False,
    initial=None,
    mask_identity=True,
    highlevel=True,
    behavior=None,
    attrs=None,
):
    """
    Args:
        array: Array-like data (anything #ak.to_layout recognizes).
        axis (None or int or str): If None, combine all values from the array into
            a single scalar result; if an int, group by that axis: `0` is the
            outermost, `1` is the first level of nested lists, etc., and
            negative `axis` counts from the innermost: `-1` is the innermost,
            `-2` is the next level up, etc; if a str, it is interpreted as the
            name of the axis which maps to an int if named axes are present.
            Named axes are attached to an array using #ak.with_named_axis and
            removed with #ak.without_named_axis; also see the
            [Named axes user guide](../../user-guide/how-to-array-properties-named-axis.html).
        keepdims (bool): If False, this reducer decreases the number of
            dimensions by 1; if True, the reduced values are wrapped in a new
            length-1 dimension so that the result of this operation may be
            broadcasted with the original array.
        initial (None or number): The minimum value of an output element, as
            an alternative to the numeric type's natural identity (e.g. negative
            infinity for floating-point types, a minimum integer for integer types).
            If you use `initial`, you might also want `mask_identity=False`.
        mask_identity (bool): If True, reducing over empty lists results in
            None (an option type); otherwise, reducing over empty lists
            results in the operation's identity.

    Like #ak.max, but treating NaN ("not a number") values as missing.

    Equivalent to

        ak.max(ak.nan_to_none(array))

    with all other arguments unchanged.

    See also #ak.max.
    """
    # Dispatch
    yield (array,)

    # Implementation
    return _impl(
        ak.operations.ak_nan_to_none._impl(array, False, None, None),
        axis,
        keepdims,
        initial,
        mask_identity,
        highlevel,
        behavior,
        attrs,
    )


def _impl(array, axis, keepdims, initial, mask_identity, highlevel, behavior, attrs):
    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
        layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")

    # Handle named axis
    named_axis = _get_named_axis(ctx)
    # Step 1: Normalize named axis to positional axis
    axis = _named_axis_to_positional_axis(named_axis, axis)
    # Step 2: propagate named axis from input to output,
    #   keepdims=True: use strategy "keep all" (see: awkward._namedaxis)
    #   keepdims=False: use strategy "remove one" (see: awkward._namedaxis)
    out_named_axis = _keep_named_axis(named_axis, None)
    if not keepdims:
        out_named_axis = _remove_named_axis(
            named_axis=out_named_axis,
            axis=axis,
            total=layout.minmax_depth[1],
        )

    axis = regularize_axis(axis, none_allowed=True)

    reducer = ak._reducers.Max(initial)

    out = ak._do.reduce(
        layout,
        reducer,
        axis=axis,
        mask=mask_identity,
        keepdims=keepdims,
        behavior=ctx.behavior,
    )

    wrapped_out = ctx.wrap(
        out,
        highlevel=highlevel,
        allow_other=True,
    )

    # propagate named axis to output
    return ak.operations.ak_with_named_axis._impl(
        wrapped_out,
        named_axis=out_named_axis,
        highlevel=highlevel,
        behavior=ctx.behavior,
        attrs=ctx.attrs,
    )


@ak._connect.numpy.implements("amax")
@ak._connect.numpy.implements("max")
def _nep_18_impl_amax(
    a, axis=None, out=UNSUPPORTED, keepdims=False, initial=None, where=UNSUPPORTED
):
    return max(a, axis=axis, keepdims=keepdims, initial=initial)


@ak._connect.numpy.implements("nanmax")
def _nep_18_impl_nanmax(
    a, axis=None, out=UNSUPPORTED, keepdims=False, initial=None, where=UNSUPPORTED
):
    return nanmax(a, axis=axis, keepdims=keepdims, initial=initial)
