import numpy as np
import xarray

from xclim.core.units import declare_units

__all__ = [
    "base_flow_index",
    "rb_flashiness_index",
]


@declare_units("", q="[discharge]")
def base_flow_index(q: xarray.DataArray, freq: str = "YS"):  # noqa: D401
    r"""Base flow index.

    Return the base flow index, defined as the minimum 7-day average flow divided by the mean flow.

    Parameters
    ----------
    q : xarray.DataArray
      Rate of river discharge [m³/s]
    freq : str
      Resampling frequency; Defaults to "YS" (yearly).

    Returns
    -------
    xarray.DataArray
      Base flow index.

    Notes
    -----
    Let :math:`\mathbf{q}=q_0, q_1, \ldots, q_n` be the sequence of daily discharge and :math:`\overline{\mathbf{q}}`
    the mean flow over the period. The base flow index is given by:

    .. math::

       \frac{\min(\mathrm{CMA}_7(\mathbf{q}))}{\overline{\mathbf{q}}}


    where :math:`\mathrm{CMA}_7` is the seven days moving average of the daily flow:

    .. math::

       \mathrm{CMA}_7(q_i) = \frac{\sum_{j=i-3}^{i+3} q_j}{7}

    """
    m7 = (
        q.rolling(time=7, center=True)
        .mean(allow_lazy=True, skipna=False)
        .resample(time=freq)
    )
    mq = q.resample(time=freq)

    m7m = m7.min(dim="time")
    return m7m / mq.mean(dim="time")


@declare_units("", q="[discharge]")
def rb_flashiness_index(q: xarray.DataArray, freq: str = "YS"):  # noqa: D401
    r"""Richards-Baker flashiness index.

    Measures oscillations in flow relative to total flow, quantifying the frequency and rapidity of short term changes
    in flow.

    Parameters
    ----------
    q : xarray.DataArray
      Rate of river discharge [m³/s]
    freq : str
      Resampling frequency; Defaults to "YS" (yearly).

    Returns
    -------
    xarray.DataArray
      R-B Index.

    Notes
    -----
    Let :math:`\mathbf{q}=q_0, q_1, \ldots, q_n` be the sequence of daily discharge, the R-B Index is given by:

    .. math::

       \frac{\sum_{i=1}^n |q_i - q_{i-1}|}{\sum_{i=1}^n q_i}

    References
    ----------
    Baker, D.B., R.P. Richards, T.T. Loftus, and J.W. Kramer, 2004. A new Flashiness Index: Characteristics and
    Applications to Midwestern Rivers and Streams. Journal of the American Water Resources Association 40(2):503-522.
    """
    d = np.abs(q.diff(dim="time")).resample(time=freq)
    mq = q.resample(time=freq)
    return d.sum(dim="time") / mq.sum(dim="time")
