"""
This module contains methods related to the Brier score
"""

import xarray as xr

from scores.continuous import mse
from scores.processing import check_binary
from scores.typing import FlexibleDimensionTypes, XarrayLike


def brier_score(
    fcst: XarrayLike,
    obs: XarrayLike,
    reduce_dims: FlexibleDimensionTypes = None,
    preserve_dims: FlexibleDimensionTypes = None,
    weights: xr.DataArray = None,
    check_args: bool = True,
):
    """
    Calculates the Brier score for forecast and observed data. For an explanation of the
    Brier score, see [Wikipedia](https://en.wikipedia.org/wiki/Brier_score).

    If you want to speed up performance with Dask and are confident of your input data,
    or if you want observations to take values between 0 and 1, set `check_inputs` to `False`.

    Args:
            fcst: Forecast or predicted variables in xarray.
            obs: Observed variables in xarray.
            reduce_dims: Optionally specify which dimensions to reduce when
                calculating the Brier score. All other dimensions will be preserved.
            preserve_dims: Optionally specify which dimensions to preserve when
                calculating the Brier 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 absolute error at each
                point (i.e. single-value comparison against observed), and the
                forecast and observed dimensions must match precisely.
            weights: Optionally provide an array for weighted averaging (e.g. by area, by latitude,
                by population, custom)
    Raises:
        ValueError: if `fcst` contains non-nan values outside of the range [0, 1]
        ValueError: if `obs` contains non-nan values not in the set {0, 1}
    """
    if check_args:
        error_msg = ValueError("`fcst` contains values outside of the range [0, 1]")
        if isinstance(fcst, xr.Dataset):
            # The .values is required as item() is not yet a valid method on Dask
            if fcst.to_array().max().values.item() > 1 or fcst.to_array().min().values.item() < 0:
                raise error_msg
        else:
            if fcst.max().values.item() > 1 or fcst.min().values.item() < 0:
                raise error_msg
        check_binary(obs, "obs")

    return mse(fcst, obs, reduce_dims, preserve_dims, weights)
