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


import awkward as ak
from awkward._attrs import attrs_of_obj
from awkward._dispatch import high_level_function
from awkward._layout import (
    HighLevelContext,
    ensure_same_backend,
    maybe_highlevel_to_lowlevel,
    promote_integral_to_float64,
)
from awkward._namedaxis import (
    AxisName,
    _get_named_axis,
)
from awkward._nplikes.numpy_like import NumpyMetadata
from awkward._typing import Mapping

__all__ = ("moment",)

np = NumpyMetadata.instance()


@high_level_function()
def moment(
    x,
    n,
    weight=None,
    axis=None,
    *,
    keepdims: bool = False,
    mask_identity: bool = False,
    highlevel: bool = True,
    behavior: Mapping | None = None,
    attrs: Mapping | None = None,
):
    """Computes the `n`th moment over one or all levels of nesting.

    Many types are supported, including all Awkward Arrays and Records. The
    grouping is performed the same way as for reducers, though this operation is
    not a reducer and has no identity.

    This function has no NumPy equivalent.

    Passing all arguments to the reducers, the moment is calculated as::

        ak.sum(x**n * weight) / ak.sum(weight)

    The `n=2` moment differs from #ak.var in that #ak.var also subtracts the
    mean (the `n=1` moment).

    See #ak.sum for a complete description of handling nested lists and
    missing values (None) in reducers, and #ak.mean for an example with another
    non-reducer.

    Args:
        x: The data on which to compute the moment (anything #ak.to_layout recognizes).
        n (int): The choice of moment: `0` is a sum of weights, `1` is
            #ak.mean, `2` is #ak.var without subtracting the mean, etc.
        weight: Data that can be broadcasted to `x` to give each value a
            weight. Weighting values equally is the same as no weights;
            weighting some values higher increases the significance of those
            values. Weights can be zero or negative.
        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 function decreases the number of
            dimensions by 1; if True, the output values are wrapped in a new
            length-1 dimension so that the result of this operation may be
            broadcasted with the original array.
        mask_identity (bool): If True, the application of this function on
            empty lists results in None (an option type); otherwise, the
            calculation is followed through with the reducers' identities,
            usually resulting in floating-point `nan`.
        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 `n`th moment in each group of elements from `x`.
    """
    # Dispatch
    yield x, weight

    # Implementation
    return _impl(
        x, n, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs
    )


def _impl(
    x,
    n,
    weight,
    axis: AxisName,
    keepdims: bool,
    mask_identity: bool,
    highlevel: bool,
    behavior: Mapping | None,
    attrs: Mapping | None,
):
    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
        x_layout, weight_layout = ensure_same_backend(
            ctx.unwrap(x, allow_record=False, primitive_policy="error"),
            ctx.unwrap(
                weight,
                allow_record=False,
                allow_unknown=False,
                primitive_policy="error",
                none_policy="pass-through",
            ),
        )

    x = ctx.wrap(x_layout)
    weight = ctx.wrap(weight_layout, allow_other=True)

    # sum-of-squares is float64-only, so complex falls back to the sum(x**n) path.
    is_complex = "complex" in str(ak.type(x))

    with np.errstate(invalid="ignore", divide="ignore"):
        if weight is None:
            sumw = ak.operations.ak_count._impl(
                x,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=ctx.behavior,
                attrs=ctx.attrs,
            )
            if is_complex:
                # The sum-of-squares/powers reducers are float64-only, so complex
                # input uses sum(x**n) (matching NumPy's complex moment).
                sumwxn = ak.operations.ak_sum._impl(
                    x**n,
                    axis,
                    keepdims,
                    mask_identity,
                    highlevel=True,
                    behavior=ctx.behavior,
                    attrs=ctx.attrs,
                )
            elif n == 2:
                # E[x**2]: sum of squares accumulated in float64 directly from
                # the input -- no x**2 buffer, no integer/float32 overflow.
                sumwxn = ak.operations.ak_sumofsquares._impl(
                    x,
                    axis,
                    keepdims,
                    mask_identity,
                    highlevel=True,
                    behavior=ctx.behavior,
                    attrs=ctx.attrs,
                )
            else:
                # Other powers: sum(x**n) accumulated in float64 directly from
                # the input -- no x**n buffer, no integer/float32 overflow.
                sumwxn = ak.operations.ak_sumofpowers._impl(
                    x,
                    n,
                    axis,
                    keepdims,
                    mask_identity,
                    highlevel=True,
                    behavior=ctx.behavior,
                    attrs=ctx.attrs,
                )
        else:
            # Promote so (x**n) * weight does not overflow for integer input.
            xp = promote_integral_to_float64(x)
            sumw = ak.operations.ak_sum._impl(
                xp * 0 + weight,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=ctx.behavior,
                attrs=ctx.attrs,
            )
            sumwxn = ak.operations.ak_sum._impl(
                (xp**n) * weight,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=ctx.behavior,
                attrs=ctx.attrs,
            )

        out = sumwxn / sumw

        # propagate named axis to output
        wrapped = ctx.wrap(
            maybe_highlevel_to_lowlevel(out),
            highlevel=highlevel,
            allow_other=True,
        )
        # propagate named axis to output
        return ak.operations.ak_with_named_axis._impl(
            wrapped,
            named_axis=_get_named_axis(attrs_of_obj(out), allow_any=True),
            highlevel=highlevel,
            behavior=None,
            attrs=None,
        )
