"""
Implementation of quantile loss (score)
"""

from typing import Optional

import xarray as xr

from scores.processing import aggregate
from scores.typing import FlexibleDimensionTypes, XarrayLike
from scores.utils import check_dims, gather_dimensions


def quantile_score(
    fcst: XarrayLike,
    obs: XarrayLike,
    alpha: float,
    *,  # Force keywords arguments to be keyword-only
    reduce_dims: Optional[FlexibleDimensionTypes] = None,
    preserve_dims: Optional[FlexibleDimensionTypes] = None,
    weights: Optional[XarrayLike] = None,
) -> XarrayLike:
    """
    Calculates a score that targets alpha-quantiles.
    Use with alpha = 0.5 for forecasts of the median.
    Use with alpha = 0.9 for forecasts of the 90th percentile.

    Args:
        fcst: array of forecasts
        obs: array of observations
        alpha: A value between 0 and 1 (exclusive)
        reduce_dims: Optionally specify which dimensions to reduce when
            calculating the quantile score. All other dimensions will be preserved. As a
            special case, 'all' will allow all dimensions to be reduced. Only one
            of `reduce_dims` and `preserve_dims` can be supplied. The default behaviour
            if neither are supplied is to reduce all dims.
        preserve_dims: Optionally specify which dimensions to preserve when calculating
            quantile score. All other dimensions will be reduced. As a special case, 'all'
            will allow all dimensions to be preserved. In this case, the result will be in
            the same shape/dimensionality as the forecast, and the errors will be the quantile
            score at each point (i.e. single-value comparison against observed), and the
            forecast and observed dimensions must match precisely. Only one of `reduce_dims`
            and `preserve_dims` can be supplied. The default behaviour if neither are supplied
            is to reduce all dims.
        weights: An array of weights to apply to the score (e.g., weighting a grid by latitude).
            If None, no weights are applied. If provided, the weights must be broadcastable
            to the data dimensions and must not contain negative or NaN values. If
            appropriate, users can choose to replace NaN values in weights by calling ``weights.fillna(0)``.
            The weighting approach follows :py:class:`xarray.computation.weighted.DataArrayWeighted`.
            See the scores weighting tutorial for more information on how to use weights.

    Returns:
        A DataArray with values being the mean generalised piecewise linear (GPL)
        scoring function, with the dimensions specified in `dims`.
        If `dims` is `None`, the returned DataArray will have only one element,
        the overall mean GPL score.

    Raises:
        ValueError: if `alpha` is not between 0 and 1.

    Notes:

        .. math::

            gpl(x) = \\begin{cases}\\alpha * (-x) & x \\leq 0\\\\
           (1-\\alpha) x & x > 0\\end{cases}

        where:
            - :math:`\\alpha` is the targeted quantile.
            - :math:`x` is the difference, fcst - obs

    References:
        T. Gneiting, "Making and evaluating point forecasts",
        J. Amer. Stat. Assoc., Vol. 106 No. 494 (June 2011), pp. 754--755,
        Theorem 9

    """
    specified_dims = reduce_dims or preserve_dims
    # check requested dims are a subset of fcst dimensions
    if specified_dims is not None and specified_dims != "all":
        check_dims(xr_data=fcst, expected_dims=specified_dims, mode="superset")
    # check obs dimensions are a subset of fcst dimensions
    check_dims(xr_data=obs, expected_dims=fcst.dims, mode="subset")
    reduce_dims = gather_dimensions(fcst.dims, obs.dims, reduce_dims=reduce_dims, preserve_dims=preserve_dims)

    # check that alpha is between 0 and 1 as required
    if (alpha <= 0) or (alpha >= 1):
        raise ValueError("alpha is not between 0 and 1")

    # Do this operation once to save compute time
    diff = fcst - obs

    # calculate the score applicable when fcst <= obs
    score_fcst_lte_obs = alpha * (-diff)

    # calculate the score applicable when fcst > obs
    score_fcst_ge_obs = (1 - alpha) * diff

    result = xr.where(diff > 0, score_fcst_ge_obs, score_fcst_lte_obs)

    score = aggregate(result, weights=weights, reduce_dims=reduce_dims)

    return score
