from collections import OrderedDict
from datetime import datetime
import pandas as pd
import numpy as np
import torch
from torch.utils.data.dataset import Dataset
from neuralprophet import hdays as hdays_part2
import holidays as hdays_part1
from collections import defaultdict
from neuralprophet import utils
import logging

log = logging.getLogger("NP.time_dataset")


class GlobalTimeDataset(Dataset):
    def __init__(self, df_dict, **kwargs):
        """Initialize Timedataset from time-series df.

        Args:
            df_dict (dict): containing pd.DataFrame time series data
            **kwargs (): identical to tabularize_univariate_datetime
        """
        self.combined_timedataset = []
        # TODO (future): vectorize
        self.length = 0
        for df_name, df in df_dict.items():
            timedataset = TimeDataset(df, df_name, **kwargs)
            self.length += timedataset.length
            for i in range(0, len(timedataset)):
                self.combined_timedataset.append(timedataset[i])

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        return self.combined_timedataset[idx]


class TimeDataset(Dataset):
    """Create a PyTorch dataset of a tabularized time-series"""

    def __init__(self, df, name, **kwargs):
        """Initialize Timedataset from time-series df.

        Args:
            df (pd.DataFrame): time series data
            name (str): name of time-series
            **kwargs (): identical to tabularize_univariate_datetime
        """
        self.name = name
        self.length = None
        self.inputs = OrderedDict({})
        self.targets = None
        self.meta = OrderedDict({})
        self.two_level_inputs = ["seasonalities", "covariates"]
        inputs, targets = tabularize_univariate_datetime(df, **kwargs)
        self.init_after_tabularized(inputs, targets)

    def init_after_tabularized(self, inputs, targets=None):
        """Create Timedataset with data.

        Args:
            inputs (ordered dict): identical to returns from tabularize_univariate_datetime
            targets (np.array, float): identical to returns from tabularize_univariate_datetime
        """
        inputs_dtype = {
            "time": torch.float,
            # "changepoints": torch.bool,
            "seasonalities": torch.float,
            "events": torch.float,
            "lags": torch.float,
            "covariates": torch.float,
            "regressors": torch.float,
        }
        targets_dtype = torch.float
        self.length = inputs["time"].shape[0]

        for key, data in inputs.items():
            if key in self.two_level_inputs or key == "events" or key == "regressors":
                self.inputs[key] = OrderedDict({})
                for name, features in data.items():
                    self.inputs[key][name] = torch.from_numpy(features).type(inputs_dtype[key])
            else:
                self.inputs[key] = torch.from_numpy(data).type(inputs_dtype[key])
        self.targets = torch.from_numpy(targets).type(targets_dtype)
        self.meta["df_name"] = self.name

    def __getitem__(self, index):
        """Overrides parent class method to get an item at index.

        Args:
            index (int): sample location in dataset

        Returns:
            sample (OrderedDict): model inputs
                time (torch tensor, float), dims: (1)
                seasonalities (OrderedDict), named seasonalities, each with features
                    (torch tensor, float) of dims: (n_features[name])
                lags (torch tensor, float), dims: (n_lags)
                covariates (OrderedDict), named covariates, each with features
                    (np.array, float) of dims: (n_lags)
                events (OrderedDict), all events both additive and multiplicative,
                    each with features (np.array, float) of dims: (n_lags)
                regressors (OrderedDict), all regressors both additive and multiplicative,
                    each with features (np.array, float) of dims: (n_lags)
            targets (torch tensor, float): targets to be predicted, dims: (n_forecasts)
        """
        # Future TODO: vectorize
        sample = OrderedDict({})
        for key, data in self.inputs.items():
            if key in self.two_level_inputs:
                sample[key] = OrderedDict({})
                for name, period_features in self.inputs[key].items():
                    sample[key][name] = period_features[index]
            elif key == "events" or key == "regressors":
                sample[key] = OrderedDict({})
                for mode, features in self.inputs[key].items():
                    sample[key][mode] = features[index, :, :]
            else:
                sample[key] = data[index]
        targets = self.targets[index]
        meta = self.meta
        return sample, targets, meta

    def __len__(self):
        """Overrides Parent class method to get data length."""
        return self.length


def tabularize_univariate_datetime(
    df,
    predict_mode=False,
    n_lags=0,
    n_forecasts=1,
    season_config=None,
    events_config=None,
    country_holidays_config=None,
    covar_config=None,
    regressors_config=None,
):
    """Create a tabular dataset from univariate timeseries for supervised forecasting.

    Note: data must be clean and have no gaps.

    Args:
        df (pd.DataFrame): Sequence of observations
            with original 'ds', 'y' and normalized 't', 'y_scaled' columns.
        season_config (configure.Season): configuration for seasonalities.
        n_lags (int): number of lagged values of series to include as model inputs. Aka AR-order
        n_forecasts (int): number of steps to forecast into future.
        events_config (OrderedDict): user specified events, each with their
            upper, lower windows (int) and regularization
        country_holidays_config (OrderedDict): Configurations (holiday_names, upper, lower windows,
            regularization) for country specific holidays
        covar_config (OrderedDict<configure.Covar>): configuration for covariates
        regressors_config (OrderedDict): configuration for regressors
        predict_mode (bool): False (default) includes target values.
            True does not include targets but includes entire dataset as input

    Returns:
        inputs (OrderedDict): model inputs, each of len(df) but with varying dimensions
            time (np.array, float), dims: (num_samples, 1)
            seasonalities (OrderedDict), named seasonalities, each with features
                (np.array, float) of dims: (num_samples, n_features[name])
            lags (np.array, float), dims: (num_samples, n_lags)
            covariates (OrderedDict), named covariates, each with features
                (np.array, float) of dims: (num_samples, n_lags)
            events (OrderedDict), events, each with features
                (np.array, float) of dims: (num_samples, n_lags)
            regressors (OrderedDict), regressors, each with features
                (np.array, float) of dims: (num_samples, n_lags)
        targets (np.array, float): targets to be predicted of same length as each of the model inputs,
            dims: (num_samples, n_forecasts)
    """
    n_samples = len(df) - n_lags + 1 - n_forecasts
    # data is stored in OrderedDict
    inputs = OrderedDict({})

    def _stride_time_features_for_forecasts(x):
        # only for case where n_lags > 0
        return np.array([x[n_lags + i : n_lags + i + n_forecasts] for i in range(n_samples)], dtype=np.float64)

    # time is the time at each forecast step
    t = df.loc[:, "t"].values
    if n_lags == 0:
        assert n_forecasts == 1
        time = np.expand_dims(t, 1)
    else:
        time = _stride_time_features_for_forecasts(t)
    inputs["time"] = time

    if season_config is not None:
        seasonalities = seasonal_features_from_dates(df["ds"], season_config)
        for name, features in seasonalities.items():
            if n_lags == 0:
                seasonalities[name] = np.expand_dims(features, axis=1)
            else:
                # stride into num_forecast at dim=1 for each sample, just like we did with time
                seasonalities[name] = _stride_time_features_for_forecasts(features)
        inputs["seasonalities"] = seasonalities

    def _stride_lagged_features(df_col_name, feature_dims):
        # only for case where n_lags > 0
        series = df.loc[:, df_col_name].values
        ## Added dtype=np.float64 to solve the problem with np.isnan for ubuntu test
        return np.array([series[i + n_lags - feature_dims : i + n_lags] for i in range(n_samples)], dtype=np.float64)

    if n_lags > 0 and "y" in df.columns:
        inputs["lags"] = _stride_lagged_features(df_col_name="y_scaled", feature_dims=n_lags)
        if np.isnan(inputs["lags"]).any():
            raise ValueError("Input lags contain NaN values in y.")

    if covar_config is not None and n_lags > 0:
        covariates = OrderedDict({})
        for covar in df.columns:
            if covar in covar_config:
                assert n_lags > 0
                window = n_lags
                if covar_config[covar].as_scalar:
                    window = 1
                covariates[covar] = _stride_lagged_features(df_col_name=covar, feature_dims=window)
                if np.isnan(covariates[covar]).any():
                    raise ValueError("Input lags contain NaN values in ", covar)

        inputs["covariates"] = covariates

    # get the regressors features
    if regressors_config is not None:
        additive_regressors, multiplicative_regressors = make_regressors_features(df, regressors_config)

        regressors = OrderedDict({})
        if n_lags == 0:
            if additive_regressors is not None:
                regressors["additive"] = np.expand_dims(additive_regressors, axis=1)
            if multiplicative_regressors is not None:
                regressors["multiplicative"] = np.expand_dims(multiplicative_regressors, axis=1)
        else:
            if additive_regressors is not None:
                additive_regressor_feature_windows = []
                for i in range(0, additive_regressors.shape[1]):
                    # stride into num_forecast at dim=1 for each sample, just like we did with time
                    additive_regressor_feature_windows.append(
                        _stride_time_features_for_forecasts(additive_regressors[:, i])
                    )
                additive_regressors = np.dstack(additive_regressor_feature_windows)
                regressors["additive"] = additive_regressors

            if multiplicative_regressors is not None:
                multiplicative_regressor_feature_windows = []
                for i in range(0, multiplicative_regressors.shape[1]):
                    # stride into num_forecast at dim=1 for each sample, just like we did with time
                    multiplicative_regressor_feature_windows.append(
                        _stride_time_features_for_forecasts(multiplicative_regressors[:, i])
                    )
                multiplicative_regressors = np.dstack(multiplicative_regressor_feature_windows)
                regressors["multiplicative"] = multiplicative_regressors

        inputs["regressors"] = regressors

    # get the events features
    if events_config is not None or country_holidays_config is not None:
        additive_events, multiplicative_events = make_events_features(df, events_config, country_holidays_config)

        events = OrderedDict({})
        if n_lags == 0:
            if additive_events is not None:
                events["additive"] = np.expand_dims(additive_events, axis=1)
            if multiplicative_events is not None:
                events["multiplicative"] = np.expand_dims(multiplicative_events, axis=1)
        else:
            if additive_events is not None:
                additive_event_feature_windows = []
                for i in range(0, additive_events.shape[1]):
                    # stride into num_forecast at dim=1 for each sample, just like we did with time
                    additive_event_feature_windows.append(_stride_time_features_for_forecasts(additive_events[:, i]))
                additive_events = np.dstack(additive_event_feature_windows)
                events["additive"] = additive_events

            if multiplicative_events is not None:
                multiplicative_event_feature_windows = []
                for i in range(0, multiplicative_events.shape[1]):
                    # stride into num_forecast at dim=1 for each sample, just like we did with time
                    multiplicative_event_feature_windows.append(
                        _stride_time_features_for_forecasts(multiplicative_events[:, i])
                    )
                multiplicative_events = np.dstack(multiplicative_event_feature_windows)
                events["multiplicative"] = multiplicative_events

        inputs["events"] = events

    if predict_mode:
        targets = np.empty_like(time)
    else:
        targets = _stride_time_features_for_forecasts(df["y_scaled"].values)

    tabularized_input_shapes_str = ""
    for key, value in inputs.items():
        if key in ["seasonalities", "covariates", "events", "regressors"]:
            for name, period_features in value.items():
                tabularized_input_shapes_str += ("    {} {} {}\n").format(name, key, period_features.shape)
        else:
            tabularized_input_shapes_str += ("    {} {} \n").format(key, value.shape)
    log.debug("Tabularized inputs shapes: \n{}".format(tabularized_input_shapes_str))

    return inputs, targets


def fourier_series(dates, period, series_order):
    """Provides Fourier series components with the specified frequency and order.

    Note: Identical to OG Prophet.

    Args:
        dates (pd.Series): containing timestamps.
        period (float): Number of days of the period.
        series_order (int): Number of fourier components.

    Returns:
        Matrix with seasonality features.
    """
    # convert to days since epoch
    t = np.array((dates - datetime(1970, 1, 1)).dt.total_seconds().astype(float)) / (3600 * 24.0)
    return fourier_series_t(t, period, series_order)


def fourier_series_t(t, period, series_order):
    """Provides Fourier series components with the specified frequency and order.

    Note: Identical to OG Prophet.

    Args:
        t (pd.Series, float): containing time as floating point number of days.
        period (float): Number of days of the period.
        series_order (int): Number of fourier components.

    Returns:
        Matrix with seasonality features.
    """
    features = np.column_stack(
        [fun((2.0 * (i + 1) * np.pi * t / period)) for i in range(series_order) for fun in (np.sin, np.cos)]
    )
    return features


def make_country_specific_holidays_df(year_list, country):
    """
    Make dataframe of country specific holidays for given years and countries

    Args:
        year_list (list): a list of years
        country (string): country name

    Returns:
        pd.DataFrame with 'ds' and 'holiday'.
    """

    try:
        country_specific_holidays = getattr(hdays_part2, country)(years=year_list)
    except AttributeError:
        try:
            country_specific_holidays = getattr(hdays_part1, country)(years=year_list)
        except AttributeError:
            raise AttributeError("Holidays in {} are not currently supported!".format(country))
    country_specific_holidays_dict = defaultdict(list)
    for date, holiday in country_specific_holidays.items():
        country_specific_holidays_dict[holiday].append(pd.to_datetime(date))
    return country_specific_holidays_dict


def make_events_features(df, events_config=None, country_holidays_config=None):
    """
    Construct arrays of all event features

    Args:
        df (pd.DataFrame): dataframe with all values including the user specified events (provided by user)
        events_config (OrderedDict): user specified events, each with their
            upper, lower windows (int), regularization
        country_holidays_config (configure.Holidays): Configurations (holiday_names, upper, lower windows, regularization)
            for country specific holidays

    Returns:
        additive_events (np.array): all additive event features (both user specified and country specific)
        multiplicative_events (np.array): all multiplicative event features (both user specified and country specific)
    """

    additive_events = pd.DataFrame()
    multiplicative_events = pd.DataFrame()

    # create all user specified events
    if events_config is not None:
        for event, configs in events_config.items():
            if event not in df.columns:
                df[event] = np.zeros_like(df["ds"], dtype=np.float64)
            feature = df[event]
            lw = configs.lower_window
            uw = configs.upper_window
            mode = configs.mode
            # create lower and upper window features
            for offset in range(lw, uw + 1):
                key = utils.create_event_names_for_offsets(event, offset)
                offset_feature = feature.shift(periods=offset, fill_value=0.0)
                if mode == "additive":
                    additive_events[key] = offset_feature
                else:
                    multiplicative_events[key] = offset_feature

    # create all country specific holidays
    if country_holidays_config is not None:
        lw = country_holidays_config.lower_window
        uw = country_holidays_config.upper_window
        mode = country_holidays_config.mode
        year_list = list({x.year for x in df.ds})
        country_holidays_dict = make_country_specific_holidays_df(year_list, country_holidays_config.country)
        for holiday in country_holidays_config.holiday_names:
            feature = pd.Series([0.0] * df.shape[0])
            if holiday in country_holidays_dict.keys():
                dates = country_holidays_dict[holiday]
                feature[df.ds.isin(dates)] = 1.0
            for offset in range(lw, uw + 1):
                key = utils.create_event_names_for_offsets(holiday, offset)
                offset_feature = feature.shift(periods=offset, fill_value=0)
                if mode == "additive":
                    additive_events[key] = offset_feature
                else:
                    multiplicative_events[key] = offset_feature

    # Make sure column order is consistent
    if not additive_events.empty:
        additive_events = additive_events[sorted(additive_events.columns.tolist())]
        additive_events = additive_events.values
    else:
        additive_events = None
    if not multiplicative_events.empty:
        multiplicative_events = multiplicative_events[sorted(multiplicative_events.columns.tolist())]
        multiplicative_events = multiplicative_events.values
    else:
        multiplicative_events = None

    return additive_events, multiplicative_events


def make_regressors_features(df, regressors_config):
    """Construct arrays of all scalar regressor features

    Args:
        df (pd.DataFrame): dataframe with all values including the user specified regressors
        regressors_config (OrderedDict): user specified regressors config

    Returns:
        additive_regressors (np.array): all additive regressor features
        multiplicative_regressors (np.array): all multiplicative regressor features

    """
    additive_regressors = pd.DataFrame()
    multiplicative_regressors = pd.DataFrame()

    for reg in df.columns:
        if reg in regressors_config:
            mode = regressors_config[reg].mode
            if mode == "additive":
                additive_regressors[reg] = df[reg]
            else:
                multiplicative_regressors[reg] = df[reg]

    if not additive_regressors.empty:
        additive_regressors = additive_regressors[sorted(additive_regressors.columns.tolist())]
        additive_regressors = additive_regressors.values
    else:
        additive_regressors = None
    if not multiplicative_regressors.empty:
        multiplicative_regressors = multiplicative_regressors[sorted(multiplicative_regressors.columns.tolist())]
        multiplicative_regressors = multiplicative_regressors.values
    else:
        multiplicative_regressors = None

    return additive_regressors, multiplicative_regressors


def seasonal_features_from_dates(dates, season_config):
    """Dataframe with seasonality features.

    Includes seasonality features, holiday features, and added regressors.

    Args:
        dates (pd.Series): with dates for computing seasonality features
        season_config (Season): configuration from NeuralProphet

    Returns:
         Dictionary with keys for each period name containing an np.array with the respective regression features.
            each with dims: (len(dates), 2*fourier_order)
    """
    assert len(dates.shape) == 1
    seasonalities = OrderedDict({})
    # Seasonality features
    for name, period in season_config.periods.items():
        if period.resolution > 0:
            if season_config.computation == "fourier":
                features = fourier_series(
                    dates=dates,
                    period=period.period,
                    series_order=period.resolution,
                )
            else:
                raise NotImplementedError
            seasonalities[name] = features
    return seasonalities
