# -*- coding: utf-8 -*-
"""
Health checks submodule
=======================

Functions performing basic health checks on xarray.DataArrays.
"""
import datetime as dt
import logging
from functools import wraps
from warnings import warn

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

from xclim.indices import generic

logging.captureWarnings(True)


# Dev notes
# ---------
#
# I suggest we use `check` for weak checking, and `assert` for strong checking.
# Weak checking would log problems in a log, while strong checking would raise an error.
#
# `functools.wraps` is used to copy the docstring and the function's original name from the source
# function to the decorated function. This allows sphinx to correctly find and document functions.


# TODO: Implement pandas infer_freq in xarray with CFTimeIndex.


def check_valid(var, key, expected):
    r"""Check that a variable's attribute has the expected value. Warn user otherwise."""

    att = getattr(var, key, None)
    if att is None:
        warn(f"Variable does not have a `{key}` attribute.", UserWarning, stacklevel=3)
    elif att != expected:
        warn(
            f"Variable has a non-conforming {key}. Got `{att}`, expected `{expected}`",
            UserWarning,
            stacklevel=3,
        )


def assert_daily(var):
    r"""Assert that the series is daily and monotonic (no jumps in time index).

    A ValueError is raised otherwise."""

    t0, t1 = var.time[:2]

    # This won't work for non-standard calendars. Needs to be implemented in xarray. Comment for now
    if isinstance(t0.values, np.datetime64):
        if pd.infer_freq(var.time.to_pandas()) != "D":
            raise ValueError("time series is not recognized as daily.")

    # Check that the first time step is one day.
    if np.timedelta64(dt.timedelta(days=1)) != (t1 - t0).data:
        raise ValueError("time series is not daily.")

    # Check that the series has the same time step throughout
    if not var.time.to_pandas().is_monotonic_increasing:
        raise ValueError("time index is not monotonically increasing.")


def check_valid_temperature(var, units):
    r"""Check that variable is air temperature."""

    check_valid(var, "standard_name", "air_temperature")
    check_valid(var, "units", units)
    assert_daily(var)


def check_valid_discharge(var):
    r"""Check that the variable is a discharge."""
    #
    check_valid(var, "standard_name", "water_volume_transport_in_river_channel")
    check_valid(var, "units", "m3 s-1")


def valid_daily_min_temperature(comp, units="K"):
    r"""Decorator to check that a computation runs on a valid temperature dataset."""

    @wraps(comp)
    def func(tasmin, *args, **kwds):
        check_valid_temperature(tasmin, units)
        check_valid(tasmin, "cell_methods", "time: minimum within days")
        return comp(tasmin, **kwds)

    return func


def valid_daily_mean_temperature(comp, units="K"):
    r"""Decorator to check that a computation runs on a valid temperature dataset."""

    @wraps(comp)
    def func(tas, *args, **kwds):
        check_valid_temperature(tas, units)
        check_valid(tas, "cell_methods", "time: mean within days")
        return comp(tas, *args, **kwds)

    return func


def valid_daily_max_temperature(comp, units="K"):
    r"""Decorator to check that a computation runs on a valid temperature dataset."""

    @wraps(comp)
    def func(tasmax, *args, **kwds):
        check_valid_temperature(tasmax, units)
        check_valid(tasmax, "cell_methods", "time: maximum within days")
        return comp(tasmax, *args, **kwds)

    return func


def valid_daily_max_min_temperature(comp, units="K"):
    r"""Decorator to check that a computation runs on valid min and max temperature datasets."""

    @wraps(comp)
    def func(tasmax, tasmin, **kwds):
        valid_daily_max_temperature(tasmax, units)
        valid_daily_min_temperature(tasmin, units)

        return comp(tasmax, tasmin, **kwds)

    return func


def valid_daily_mean_discharge(comp):
    r"""Decorator to check that a computation runs on valid discharge data."""

    @wraps(comp)
    def func(q, **kwds):
        check_valid_discharge(q)
        return comp(q, **kwds)

    return func


def valid_missing_data_threshold(comp, threshold=0):
    r"""Check that the relative number of missing data points does not exceed a threshold."""
    # TODO
    raise NotImplementedError


def check_is_dataarray(comp):
    r"""Decorator to check that a computation has an instance of xarray.DataArray
     as first argument."""

    @wraps(comp)
    def func(data_array, *args, **kwds):
        assert isinstance(data_array, xr.DataArray)
        return comp(data_array, *args, **kwds)

    return func


# This function can probably be made simpler once CFPeriodIndex is implemented.
def missing_any(da, freq, **indexer):
    r"""Return a boolean DataArray indicating whether there are missing days in the resampled array.

    Parameters
    ----------
    da : DataArray
      Input array at daily frequency.
    freq : str
      Resampling frequency.
    **indexer : {dim: indexer, }, optional
      Time attribute and values over which to subset the array. For example, use season='DJF' to select winter values,
      month=1 to select January, or month=[6,7,8] to select summer months. If not indexer is given, all values are
      considered.

    Returns
    -------
    out : DataArray
      A boolean array set to True if any month or year has missing values.
    """
    if "-" in freq:
        pfreq, anchor = freq.split("-")
    else:
        pfreq = freq

    # Compute the number of days in the time series during each period at the given frequency.
    selected = generic.select_time(da, **indexer)
    if selected.time.size == 0:
        raise ValueError("No data for selected period.")

    c = selected.notnull().resample(time=freq).sum(dim="time")

    # Otherwise simply use the start and end dates to find the expected number of days.
    if pfreq.endswith("S"):
        start_time = c.indexes["time"]
        end_time = start_time.shift(1, freq=freq)
    else:
        end_time = c.indexes["time"]
        start_time = end_time.shift(-1, freq=freq)

    if indexer:
        # Create a full synthetic time series and compare the number of days with the original series.
        t0 = str(start_time[0].date())
        t1 = str(end_time[-1].date())
        if isinstance(c.indexes["time"], xr.CFTimeIndex):
            cal = da.time.encoding.get("calendar")
            t = xr.cftime_range(t0, t1, freq="D", calendar=cal)
        else:
            t = pd.date_range(t0, t1, freq="D")

        sda = xr.DataArray(data=np.ones(len(t)), coords={"time": t}, dims=("time",))
        st = generic.select_time(sda, **indexer)
        sn = st.notnull().resample(time=freq).sum(dim="time")
        miss = sn != c
        return miss

    n = (end_time - start_time).days
    nda = xr.DataArray(n.values, coords={"time": c.time}, dims="time")
    return c != nda
