# 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 _get_named_axis, _is_valid_named_axis
from awkward._nplikes.numpy_like import NumpyMetadata
from awkward._regularize import regularize_axis

__all__ = ("covar",)

np = NumpyMetadata.instance()


@high_level_function()
def covar(
    x,
    y,
    weight=None,
    axis=None,
    *,
    keepdims=False,
    mask_identity=False,
    highlevel=True,
    behavior=None,
    attrs=None,
):
    """Computes the covariance of x and y over one or all levels of nesting.

    Many types are supported, including all Awkward Arrays and Records, which
    must be broadcastable to each other. 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 covariance is calculated as::

        ak.sum((x - ak.mean(x))*(y - ak.mean(y))*weight) / ak.sum(weight)

    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: One coordinate to use in the covariance calculation (anything #ak.to_layout recognizes).
        y: The other coordinate to use in the covariance calculation (anything #ak.to_layout recognizes).
        weight: Data that can be broadcasted to `x` and `y` to give each point
            a weight. Weighting points equally is the same as no weights;
            weighting some points higher increases the significance of those
            points. 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 covariance of `x` and `y`.
    """
    # Dispatch
    yield x, y, weight

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


def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs):
    if _is_valid_named_axis(axis):
        raise NotImplementedError("named axis not yet supported for ak.covar")
    axis = regularize_axis(axis, none_allowed=True)

    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
        x_layout, y_layout, weight_layout = ensure_same_backend(
            ctx.unwrap(x, allow_record=False, primitive_policy="error"),
            ctx.unwrap(y, 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)
    y = ctx.wrap(y_layout)
    weight = ctx.wrap(weight_layout, allow_other=True)

    with np.errstate(invalid="ignore", divide="ignore"):
        xmean = ak.operations.ak_mean._impl(
            x,
            weight,
            axis,
            True,
            mask_identity,
            highlevel=True,
            behavior=None,
            attrs=None,
        )
        ymean = ak.operations.ak_mean._impl(
            y,
            weight,
            axis,
            True,
            mask_identity,
            highlevel=True,
            behavior=None,
            attrs=None,
        )
        if weight is None:
            sumw = ak.operations.ak_count._impl(
                x,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=None,
                attrs=None,
            )
        else:
            sumw = ak.operations.ak_sum._impl(
                x * 0 + weight,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=None,
                attrs=None,
            )

        # Centring is internal; strip named axes from the means so the
        # subtractions are not rejected by the named-axis compatibility check.
        xmean = ak.operations.ak_without_named_axis._impl(
            xmean, highlevel=True, behavior=None, attrs=None
        )
        ymean = ak.operations.ak_without_named_axis._impl(
            ymean, highlevel=True, behavior=None, attrs=None
        )

        # Centring on the (float64) means -- numerically stable. For axis=None the
        # mean is a scalar, so subtract it directly; broadcasting the length-1
        # keepdims mean against a ragged array instead is far slower and heavier.
        # Otherwise two-pass, falling back to the one-pass E[xy]-E[x]E[y] form when
        # `x - xmean` cannot broadcast (a non-innermost axis over a ragged array,
        # or an empty axis=None).
        if axis is None:
            xm = xmean[(0,) * xmean.ndim]
            ym = ymean[(0,) * ymean.ndim]
            prod = None if (xm is None or ym is None) else (x - xm) * (y - ym)
        else:
            try:
                prod = (x - xmean) * (y - ymean)
            except ValueError:
                prod = None

        if prod is not None:
            if weight is not None:
                prod = prod * weight
            sumwxy = ak.operations.ak_sum._impl(
                prod,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=None,
                attrs=None,
            )
            out = sumwxy / sumw
        else:
            xp = promote_integral_to_float64(x)
            yp = promote_integral_to_float64(y)
            xy = xp * yp if weight is None else xp * yp * weight
            sumwxy = ak.operations.ak_sum._impl(
                xy,
                axis,
                keepdims,
                mask_identity,
                highlevel=True,
                behavior=None,
                attrs=None,
            )

            # xmean and ymean were forced to keepdims=True earlier for the try block.
            # If keepdims=False, we must recalculate them so they match the dimensions of sumwxy.
            if not keepdims:
                xmean = ak.operations.ak_mean._impl(
                    x,
                    weight,
                    axis,
                    False,
                    mask_identity,
                    highlevel=True,
                    behavior=None,
                    attrs=None,
                )
                ymean = ak.operations.ak_mean._impl(
                    y,
                    weight,
                    axis,
                    False,
                    mask_identity,
                    highlevel=True,
                    behavior=None,
                    attrs=None,
                )
                xmean = ak.operations.ak_without_named_axis._impl(
                    xmean, highlevel=True, behavior=None, attrs=None
                )
                ymean = ak.operations.ak_without_named_axis._impl(
                    ymean, highlevel=True, behavior=None, attrs=None
                )

            out = sumwxy / sumw - xmean * ymean

        wrapped = ctx.wrap(
            maybe_highlevel_to_lowlevel(out),
            highlevel=highlevel,
            allow_other=True,
        )
        return ak.operations.ak_with_named_axis._impl(
            wrapped,
            named_axis=_get_named_axis(attrs_of_obj(out)),
            highlevel=highlevel,
            behavior=None,
            attrs=None,
        )
