"""
This module contains a variety of functions which modify data in various ways to process data structures
to support probablistic verification.
"""

from typing import Iterable, Literal, Optional

import numpy as np
import pandas as pd
import xarray as xr

from .checks import cdf_values_within_bounds, check_nan_decreasing_inputs


def round_values(array: xr.DataArray, rounding_precision: float, final_round_decpl: int = 7) -> xr.DataArray:
    """
    Round data array to specified precision.
    Assumes that rounding_precision >=0, with 0 indicating no rounding to be performed.
    For example, 3.73 rounded to precision 0.2 is 3.8.
    If rounding_precision > 0, a final round to `final_round_decpl` decimal places is performed
    to remove artefacts of python rounding process.

    Args:
        array: array of data to be rounded
        rounding_precision: rounding precision
        final_round_decpl: final round to specified number of decimal places when `rounding_precision` > 0.

    Returns:
        An xarray data array with rounded values.

    Raises:
        ValueError: If `rounding_precision` < 0.
    """
    if rounding_precision < 0:
        raise ValueError(f"rounding_precision '{rounding_precision}' is negative")

    if rounding_precision > 0:
        array = (array / rounding_precision).round() * rounding_precision
        array = array.round(decimals=final_round_decpl)

    return array


def propagate_nan(cdf: xr.DataArray, threshold_dim: str) -> xr.DataArray:
    """
    If there is a NaN value in xr.DataArray `cdf`, then `propagate_nan` propagates the NaN value
    along the `threshold_dim` dimension in both directions.

    Args:
        cdf: DataArray of CDF values, so that P(X <= threshold) = cdf_value for each threshold
            in the `threshold_dim` dimension.
        threshold_dim: name of the threshold dimension in `cdf`.

    Returns:
        An xarray DataArray of `cdf` but with NaNs propagated.

    Raises:
        ValueError: If `threshold_dim` is not a dimension of `cdf`.
    """
    if threshold_dim not in cdf.dims:
        raise ValueError(f"'{threshold_dim}' is not a dimension of `cdf`")

    where_nan = np.isnan(cdf).any(dim=threshold_dim)
    result = cdf.where(~where_nan, np.nan)
    return result


def observed_cdf(
    obs: xr.DataArray,
    threshold_dim: str,
    threshold_values: Optional[Iterable[float]] = None,
    include_obs_in_thresholds: bool = True,
    precision: float = 0,
) -> xr.DataArray:
    """
    Returns a data array of observations converted into CDF format, such that
        returned_value = 0 if threshold < observation
        returned_value = 1 if threshold >= observation

    Args:
        obs: observations
        threshold_dim: name of dimension in returned array that contains the threshold values.
        threshold_values: values to include among thresholds.
        include_obs_in_thresholds: if `True`, include (rounded) observed values among thresholds.
        precision: precision applied to observed values prior to constructing the CDF and
            thresholds. Select 0 for highest precision (i.e. no rounding).

    Returns:
        an xarray DataArray with observed CDFs and thresholds in the `threshold_dim` dimension.

    Raises:
        ValueError: if `precision < 0`.
        ValueError: if all observations are NaN and no non-NaN `threshold_values`
            are not supplied.

    """
    if precision < 0:
        raise ValueError("`precision` must be nonnegative.")

    if np.isnan(obs).all() and (threshold_values is None or np.isnan(threshold_values).all()):
        raise ValueError("must include non-NaN observations in thresholds or supply threshold values")

    if precision > 0:
        obs = round_values(obs, precision)

    thresholds = threshold_values if threshold_values is not None else []

    if include_obs_in_thresholds:
        thresholds = np.concatenate((obs.values.flatten(), thresholds))

    # remove any NaN
    thresholds = [x for x in thresholds if not np.isnan(x)]

    thresholds = np.sort(pd.unique(thresholds))

    da_thresholds = xr.DataArray(
        thresholds,
        dims=[threshold_dim],
        coords={threshold_dim: thresholds},
    )

    dab_obs, dab_thresholds = xr.broadcast(obs, da_thresholds)

    cdf = dab_thresholds >= dab_obs

    # convert to 0 and 1
    cdf = cdf.astype(float)

    cdf = cdf.where(~np.isnan(dab_obs))

    return cdf


def integrate_square_piecewise_linear(function_values: xr.DataArray, threshold_dim: str) -> xr.DataArray:
    """
    Calculates integral(F(t)^2), where
        - If t in a threshold value in `threshold_dim` then F(t) is in `function_values`,
        - F is piecewise linear between each of the t values in `threshold_dim`.
    Returns value of the integral with `threshold_dim` collapsed and other dimensions preserved.
    Returns NaN if there are less than two non-NaN function_values.

    This function assumes that
        - `threshold_dim` is a dimension of `function_values`
        - coordinates of `threshold_dim` are increasing.

    Args:
        function_values: array of function values F(t).
        threshold_dim: dimension along which to integrate.

    Returns:
        An xarray DataArray of integral values and `threshold_dim` collapsed.
    """

    # notation: Since F is piecewise linear we have
    # F(t) = mt + b, whenever x[i-1] <= t <= x[i].

    # difference in x
    diff_xs = function_values[threshold_dim] - function_values[threshold_dim].shift(**{threshold_dim: 1})

    # difference in function values
    diff_ys = function_values - function_values.shift(**{threshold_dim: 1})

    # gradients m
    m_values = diff_ys / diff_xs

    # y intercepts b
    b_values = function_values.shift(**{threshold_dim: 1})

    # integral for x[i-1] <= t <= x[i]
    piece_integral = (
        (m_values**2) * (diff_xs**3) / 3 + m_values * b_values * (diff_xs**2) + (b_values**2) * diff_xs
    )

    # Need at least one non-NaN piece_integral to return float.
    # Note: need at least two non-NaN function values to get a non-NaN piece_integral.
    return piece_integral.sum(threshold_dim, min_count=1)


def add_thresholds(
    cdf: xr.DataArray,
    threshold_dim: str,
    new_thresholds: Iterable[float],
    fill_method: Literal["linear", "step", "forward", "backward", "none"],
    min_nonnan: int = 2,
) -> xr.DataArray:
    """
    Takes a CDF data array with dimension `threshold_dim` and adds values from
    `new_thresholds` to the `threshold_dim` dimension.
    The CDF is then filled to replace any NaN values.
    The array `cdf` requires at least 2 non-NaN values along `threshold_dim`.

    Args:
        cdf: array of CDF values.
        threshold_dim: name of the threshold dimension in `cdf`.
        new_thresholds: new thresholds to add to `cdf`.
        fill_method: one of "linear", "step", "forward" or "backward", as described in `fill_cdf`.
            If no filling, set to "none".
        min_nonnan: passed onto `fill_cdf` for performing filling.

    Returns:
        an xarray DataArray with additional thresholds, and values at those thresholds
        determined by the specified fill method.
    """

    thresholds = np.concatenate((cdf[threshold_dim].values, new_thresholds))
    thresholds = np.sort(pd.unique(thresholds))
    thresholds = thresholds[~np.isnan(thresholds)]

    da_thresholds = xr.DataArray(data=thresholds, dims=[threshold_dim], coords={threshold_dim: thresholds})

    da_cdf = xr.broadcast(cdf, da_thresholds)[0]

    if fill_method != "none":
        da_cdf = fill_cdf(da_cdf, threshold_dim, fill_method, min_nonnan)

    return da_cdf


def fill_cdf(
    cdf: xr.DataArray,
    threshold_dim: str,
    method: Literal["linear", "step", "forward", "backward"],
    min_nonnan: int,
) -> xr.DataArray:
    """
    Fills NaNs in a CDF of a real-valued random variable along `threshold_dim`
    with appropriate values between 0 and 1.

    Args:
        cdf: CDF values, where P(Y <= threshold) = cdf_value for each threshold in `threshold_dim`.
        threshold_dim: the threshold dimension in the CDF, along which filling is performed.
        method: one of
            - "linear": use linear interpolation, and if needed also extrapolate linearly. Clip to 0 and 1.
                Needs at least two non-NaN values for interpolation, so returns NaNs where this condition fails.
            - "step": use forward filling then set remaining leading NaNs to 0.
                Produces a step function CDF (i.e. piecewise constant).
            - "forward": use forward filling then fill any remaining leading NaNs with backward filling.
            - "backward": use backward filling then fill any remaining trailing NaNs with forward filling.
        min_nonnan: the minimum number of non-NaN entries required along `threshold_dim` for filling to
            be performed. All CDF values are set to `np.nan` where this condition fails.
            `min_nonnan` must be at least 2 for the "linear" method, and at least 1 for the other methods.

    Returns:
        An xarray DataArray with the same values as `cdf` but with NaNs filled.

    Raises:
        ValueError: If `threshold_dim` is not a dimension of `cdf`.
        ValueError: If `min_nonnan` < 1 when `method="step"`
            or if `min_nonnan` < 2 when `method="linear"`.
        ValueError: If `method` is not "linear", "step", "forward" or "backward".
        ValueError: If any non-NaN value of `cdf` lies outside the unit interval [0,1].

    """

    if method not in ["linear", "step", "forward", "backward"]:
        raise ValueError("`method` must be 'linear', 'step', 'forward' or 'backward'")

    if not cdf_values_within_bounds(cdf):
        raise ValueError("Input CDF has some values less than 0 or greater than 1.")

    if threshold_dim not in cdf.dims:
        raise ValueError(f"'{threshold_dim}' is not a dimension of `cdf`")

    if min_nonnan < 1 and method != "linear":
        raise ValueError(f"`min_nonnan` must be at least 1 when `method='{method}'`")

    if min_nonnan < 2 and method == "linear":
        raise ValueError("`min_nonnan` must be at least 2 when `method='linear'`")

    # set cdf values to be NaN where min_nonnan requirement fails
    nonnan_count = cdf.count(threshold_dim)
    cdf = cdf.where(nonnan_count >= min_nonnan)

    # NaN filling
    if method == "linear":
        cdf = cdf.interpolate_na(threshold_dim, method="linear", fill_value="extrapolate").clip(min=0, max=1)

    if method == "step":
        cdf = cdf.ffill(threshold_dim).fillna(0)
        # NaN cdfs will now be all zero, so bring back Nans
        cdf = cdf.where(nonnan_count >= min_nonnan)

    if method == "forward":
        cdf = cdf.ffill(threshold_dim).bfill(threshold_dim)

    if method == "backward":
        cdf = cdf.bfill(threshold_dim).ffill(threshold_dim)

    return cdf


def decreasing_cdfs(cdf: xr.DataArray, threshold_dim: str, tolerance: float) -> xr.DataArray:
    """
    A CDF of a real-valued random variable should be nondecreasing along threshold_dim.
    This is sometimes violated due to rounding issues or bad forecast process.
    `decreasing_cdfs` checks CDF values decrease beyond specified tolerance; that is,
    whenever the sum of the incremental decreases exceeds tolerarance.
    For example, if the CDF values are
        [0, 0.4, 0.3, 0.9, 0.88, 1]
    then the sum of incremental decreases is -0.12. Given a specified positive `tolerance`,
    the CDF values decrease beyond tolerance if the sum of incremental decreases < -`tolerance`.

    Intended use is for CDFs with increasing coordinates along `threshold_dim` dimension, and where
    either each CDF is always NaN or always non-NaN.

    Args:
        cdf: data array of CDF values
        threshold_dim: threshold dimension, such that P(Y < threshold) = cdf_value.
        tolerance: nonnegative tolerance value.

    Returns:
        An xarray DataArray with `threshold_dim` collapsed and values True if and only if
        the CDF is decreasing outside tolerance. If the CDF consists only of NaNs then
        the value is False.


    Raises:
        ValueError: If `threshold_dim` is not a dimension of `cdf`.
        ValueError: If `tolerance` is negative.
        ValueError: If coordinates are not increasing along `threshold_dim`.
        ValueError: If some, but not all, CDF values in `cdf` along `threshold_dim` are NaN.
    """
    check_nan_decreasing_inputs(cdf, threshold_dim, tolerance)

    # difference between consecutive terms along threshold_dim
    diff = cdf - cdf.shift(**{threshold_dim: 1})

    result = diff.clip(max=0).sum(dim=threshold_dim) < -tolerance

    return result


def cdf_envelope(
    cdf: xr.DataArray,
    threshold_dim: str,
) -> xr.DataArray:
    """
    Forecast cumulative distribution functions (CDFs) for real-valued random variables
    that are reconstructed from known points on the
    distribution should be nondecreasing with respect to the threshold dimension.
    However, sometimes this may fail due to rounding or poor forecast process.
    This function returns the "envelope" of the original CDF, which consists of
    two bounding CDFs, both of which are nondecreasing.

    The following example shows values from an original CDF that has a decreasing subsequence
    (and so is not a true CDF). The resulting "upper" and "lower" CDFs minimally adjust
    "original" so that "lower" <= "original" <= "upper".
        "original": [0, .5, .2, .8, 1]
        "upper": [0, .5, .5, .8, 1]
        "lower": [0, .2, .2, .8, 1]

    This function does not perform checks that `0 <= cdf <= 1`.

    Args:
        cdf: forecast CDF with thresholds in the thresholds_dim.
        threshold_dim: dimension in fcst_cdf that contains the threshold ordinates.

    Returns:
        An xarray DataArray consisting of three CDF arrays indexed along the `"cdf_type"` dimension
        with the following indices:
            - "original": same data as `cdf`.
            - "upper": minimally adjusted "original" CDF that is nondecreasing and
                satisfies "upper" >= "original".
            - "lower": minimally adjusted "original" CDF that is nondecreasing and
                satisfies "lower" <= "original".
        NaN values in `cdf` are maintained in "original", "upper" and "lower".

    Raises:
        ValueError: If `threshold_dim` is not a dimension of `cdf`.
    """
    if threshold_dim not in cdf.dims:
        raise ValueError(f"'{threshold_dim}' is not a dimension of `cdf`")

    # logic below assumes that cdf[threshold_dim].values are ascending
    cdf = cdf.sortby(threshold_dim)

    result = xr.full_like(cdf.expand_dims({"cdf_type": ["original", "upper", "lower"]}), np.nan)

    dim_idx = cdf.dims.index(threshold_dim)

    # use fmax so as not to propogate nans
    cdf_upper = np.fmax.accumulate(cdf.values, axis=dim_idx)
    cdf_lower = np.flip(
        1 - np.fmax.accumulate(1 - np.flip(cdf.values, axis=dim_idx), axis=dim_idx),
        axis=dim_idx,
    )

    result.loc["original"] = cdf.copy()
    result.loc["upper"] = np.where(~np.isnan(cdf), cdf_upper, np.nan)
    result.loc["lower"] = np.where(~np.isnan(cdf), cdf_lower, np.nan)

    return result
