"""Commonly used utility functions for Lilio."""
import re
import typing
import warnings
from typing import Dict
from typing import List
from typing import Union
import numpy as np
import pandas as pd
import xarray as xr


if typing.TYPE_CHECKING:
    from lilio import Calendar


MONTH_LENGTH = 30  # Month length for Timedelta checks.


def check_timeseries(
    data: Union[pd.Series, pd.DataFrame, xr.DataArray, xr.Dataset]
) -> None:
    """Check if input data contains valid time data.

    Checks if:
     - Input data is pd.Dataframe/pd.Series/xr.Dataset/xr.DataArray.
     - Input data has a time index (pd), or a dim named `time` containing datetime
       values
    """
    if not isinstance(data, (pd.Series, pd.DataFrame, xr.DataArray, xr.Dataset)):
        raise ValueError("The input data is neither a pandas or xarray object")
    if isinstance(data, (pd.Series, pd.DataFrame)):
        check_time_dim_pandas(data)
    elif isinstance(data, (xr.DataArray, xr.Dataset)):
        check_time_dim_xarray(data)


def is_dask_array(data: Union[xr.DataArray, xr.Dataset]) -> bool:
    """Checks if the xarray dataset/array has any dask arrays."""
    if isinstance(data, xr.DataArray):
        return False if data.chunks is None else True

    if any([data[var].chunks is not None for var in list(data.variables)]):
        return True
    else:
        return False


def check_time_dim_xarray(data) -> None:
    """Check if an xarray data has a time dimensions with time data."""
    if "time" not in data.dims:
        raise ValueError(
            "The input DataArray/Dataset does not contain a `time` dimension"
        )
    if not xr.core.common.is_np_datetime_like(data["time"].dtype):  # type: ignore
        raise ValueError("The `time` dimension is not of a datetime format")


def check_time_dim_pandas(data) -> None:
    """Check if pandas data has an index with time data."""
    if not isinstance(data.index, pd.DatetimeIndex):
        raise ValueError("The input data does not have a datetime index.")


def check_empty_intervals(indices_list: List[np.ndarray]) -> None:
    """Check for empty intervals in the resampling data.

    Args:
        indices_list: A list, where each item is an array with indices corresponding
            to the to-be-resampled data's time axis.

    Raises:
        UserWarning: If the data is insufficient.

    Returns:
        None
    """
    if any(len(interval_locs) == 1 for interval_locs in indices_list):
        warnings.warn(  # type: ignore
            message=(
                "\n  Some intervals only contains a single data point."
                "\n  This could lead to issues like aliasing or incorrect resampling."
                "\n  If possible: make the Calendar's intervals larger, or use data of"
                "\n  a higher time resolution."
            ),
            stacklevel=1,
        )
    elif any(len(interval_locs) == 0 for interval_locs in indices_list):
        warnings.warn(  # type: ignore
            message=(
                "\n  The input data could not fully cover the calendar's intervals."
                "\n  Intervals without available data will contain NaN values.",
                "\n  If possible: make the Calendar's intervals larger, or use data of"
                "\n  a higher time resolution.",
            ),
            stacklevel=1,
        )


def infer_input_data_freq(
    data: Union[pd.Series, pd.DataFrame, xr.DataArray, xr.Dataset]
) -> pd.Timedelta:
    """Infer the frequency of the input data, for comparison with the calendar freq.

    Args:
        data: Pandas or xarray data to infer the frequency of.

    Returns:
        a pd.Timedelta
    """
    if isinstance(data, (pd.Series, pd.DataFrame)):
        data_freq = pd.infer_freq(data.index)
        if data_freq is None:  # Manually infer the frequency
            data_freq = np.min(data.index.values[1:] - data.index.values[:-1])
    else:
        data_freq = xr.infer_freq(data.time)
        if data_freq is None:  # Manually infer the frequency
            data_freq = (data.time.values[1:] - data.time.values[:-1]).min()

    if isinstance(data_freq, str):
        data_freq.replace("-", "")  # Get the absolute frequency

        if not re.match(r"\d+\D", data_freq):  # infer_freq can return "d" for "1d".
            data_freq = "1" + data_freq

        data_freq = (  # Deal with monthly timedelta case
            replace_month_length(data_freq) if data_freq[-1] == "M" else data_freq
        )
    return pd.Timedelta(data_freq)


def replace_month_length(length: str) -> str:
    """Replace month lengths with an equivalent length in days."""
    ndays = int(length[:-1]) * MONTH_LENGTH
    return f"{ndays}d"


def get_smallest_calendar_freq(calendar: "Calendar") -> pd.Timedelta:
    """Return the smallest length of the calendar's intervals as a Timedelta."""
    intervals = calendar.targets + calendar.precursors
    lengthstr = [iv.length for iv in intervals]
    lengthstr = [ln.replace("-", "") for ln in lengthstr]  # Account for neg. lengths
    lengthstr = [replace_month_length(ln) if ln[-1] == "M" else ln for ln in lengthstr]
    lengths = [pd.Timedelta(ln) for ln in lengthstr]
    return min(lengths)


def check_input_frequency(
    calendar: "Calendar", data: Union[pd.Series, pd.DataFrame, xr.DataArray, xr.Dataset]
):
    """Compare the frequency of (input) data to the frequency of the calendar.

    Note: Pandas and xarray have the builtin function `infer_freq`, but this function is
    not robust enough for our purpose, so we have to manually infer the frequency if the
    builtin one fails.
    """
    data_freq = infer_input_data_freq(data)
    calendar_freq = get_smallest_calendar_freq(calendar)

    if calendar_freq < data_freq:
        raise ValueError(
            "The data is of a lower time resolution than the calendar. "
            "This would lead to incorrect data and/or NaN values in the resampled data."
            " Please make the Calendar's intervals larger, or use data of a higher time"
            " resolution."
            f"\nInfered data frequency: {str(data_freq)} < calendar frequency "
            f"{str(calendar_freq)}"
        )
    if calendar_freq < 2 * data_freq:
        warnings.warn(
            "\n  The input data frequency is very close to the Calendar's frequency."
            "\n  This could lead to issues like aliasing or incorrect resampling."
            "\n  If possible: make the Calendar's intervals larger, or use data of a"
            "\n  higher time resolution.",
            stacklevel=1,
        )


def convert_interval_to_bounds(data: xr.Dataset) -> xr.Dataset:
    """Convert pandas intervals to bounds in a xarray Dataset.

    pd.Interval objects cannot be written to netCDF. To allow writing the
    calendar-resampled data to netCDF these intervals have to be converted to bounds.
    This function adds a 'bounds' dimension, with 'left' and 'right' coordinates, and
    converts the 'interval' coordinates to this system.

    Args:
        data: Input data with intervals as pd.Interval objects.

    Returns:
        Input data with the intervals converted to bounds.
    """
    data = data.stack(coord=["anchor_year", "i_interval"])
    bounds = np.array([[val.left, val.right] for val in data.interval.values])
    data["left_bound"] = ("coord", bounds[:, 0])
    data["right_bound"] = ("coord", bounds[:, 1])
    data["left_bound"].attrs = {
        "name": "Left bound of the interval",
        "closed": "True",
    }
    data["right_bound"].attrs = {
        "name": "Right bound of the interval",
        "closed": "False",
    }
    data = data.unstack("coord")
    data = data.drop_vars(["interval"])
    data = data.set_coords(["left_bound", "right_bound"])
    return data


def check_reserved_names(
    input_data: Union[pd.Series, pd.DataFrame, xr.DataArray, xr.Dataset]
) -> None:
    """Check if reserved names are already in the input data. E.g. "anchor_year"."""
    reserved_names_pd = ["anchor_year", "i_interval", "is_target"]
    reserved_names_xr = reserved_names_pd + ["left_bound", "right_bound"]

    if isinstance(input_data, pd.DataFrame):
        if any(name in input_data.columns for name in reserved_names_pd):
            raise ValueError(
                "The input data contains one or more reserved names. Please remove or "
                f"rename these before continuing.\n Reserved names: {reserved_names_pd}"
            )
    elif isinstance(input_data, (xr.DataArray, xr.Dataset)):
        data_names = [
            input_data.keys()
            if isinstance(input_data, xr.Dataset)
            else list(input_data.coords) + [input_data.name]
        ]
        if any(name in data_names for name in reserved_names_xr):
            raise ValueError(
                "The input data contains one or more reserved names. Please remove or "
                f"rename these before continuing.\n Reserved names: {reserved_names_xr}"
            )


def assert_bokeh_available():
    """Util that attempts to load the optional module bokeh."""
    try:
        import bokeh as _  # noqa: F401 (unused import)

    except ImportError as e:
        raise ImportError(
            "Could not import the `bokeh` module.\nPlease install this"
            " before continuing, with either `pip` or `conda`."
        ) from e


def get_month_names() -> Dict:
    """Generate a dictionary with English lowercase month names and abbreviations.

    Returns:
        Dictionary containing the English names of the months, including their
            abbreviations, linked to the number of each month.
            E.g. {'december': 12, 'jan': 1}
    """
    return {
        "january": 1,
        "february": 2,
        "march": 3,
        "april": 4,
        "may": 5,
        "june": 6,
        "july": 7,
        "august": 8,
        "september": 9,
        "october": 10,
        "november": 11,
        "december": 12,
        "jan": 1,
        "feb": 2,
        "mar": 3,
        "apr": 4,
        "jun": 6,
        "jul": 7,
        "aug": 8,
        "sep": 9,
        "oct": 10,
        "nov": 11,
        "dec": 12,
    }


def check_month_day(month: int, day: int = 1):
    """Check if the input day/month combination is valid.

    Months must be between 1 and 12, and days must be within 1 and 28/30/31 (depending
    on the month).

    Args:
        month: Month number
        day: Day number. Defaults to 1.
    """
    if month in {1, 3, 5, 7, 8, 10, 12}:
        if (day < 1) or (day > 31):
            raise ValueError(
                "Incorrect anchor input. "
                f"Day number {day} is not a valid day for month {month}"
            )
    elif month in {4, 6, 9, 11}:
        if (day < 1) or (day > 30):
            raise ValueError(
                "Incorrect anchor input. "
                f"Day number {day} is not a valid day for month {month}"
            )
    elif month == 2:
        if (day < 1) or (day > 28):
            raise ValueError(
                "Incorrect anchor input. "
                f"Day number {day} is not a valid day for month {month}"
            )
    else:
        raise ValueError(
            "Incorrect anchor input. Month number must be between 1 and 12."
        )


def check_week_day(week: int, day: int = 1):
    """Validate the week and day numbers."""
    if week == 53:
        raise ValueError(
            "Incorrect anchor input. "
            "Week 53 is not a valid input, as not every year contains a 53rd week."
        )
    if (week < 1) or (week > 52):
        raise ValueError(
            "Incorrect anchor input. Week numbers must be between 1 and 52."
        )
    if (day < 1) or (day > 7):
        raise ValueError(
            "Incorrect anchor input. Weekday numbers must be between 1 and 7."
        )


def parse_freqstr_to_dateoffset(time_str):
    """Parse the user-input time strings.

    Args:
        time_str: Time length string in the right formatting.

    Returns:
        Dictionary as keyword argument for Pandas DateOffset.
    """
    if re.fullmatch(r"[+-]?\d*d", time_str):
        time_dict = {"days": int(time_str[:-1])}
    elif re.fullmatch(r"[+-]?\d*M", time_str):
        time_dict = {"months": int(time_str[:-1])}
    elif re.fullmatch(r"[+-]?\d*W", time_str):
        time_dict = {"weeks": int(time_str[:-1])}
    else:
        raise ValueError("Please input a time string in the correct format.")

    return time_dict
