# Authors: Thierry Moudiki
#
# License: BSD 3

import numpy as np
from ..base import Base
from scipy.stats import norm
from ..utils import matrixops as mo
from ..utils import timeseries as ts
import sklearn.metrics as skm2
import pickle
from functools import partial


class MTS(Base):
    """Univariate and multivariate time series (MTS) forecasting with Quasi-Randomized networks
    
    Attributes: 
     
        obj: object.
            any object containing a method fit (obj.fit()) and a method predict 
            (obj.predict()).

        n_hidden_features: int.
            number of nodes in the hidden layer.

        activation_name: str.
            activation function: 'relu', 'tanh', 'sigmoid', 'prelu' or 'elu'.

        a: float.
            hyperparameter for 'prelu' or 'elu' activation function.

        nodes_sim: str.
            type of simulation for the nodes: 'sobol', 'hammersley', 'halton', 
            'uniform'.

        bias: boolean.
            indicates if the hidden layer contains a bias term (True) or not 
            (False).

        dropout: float.
            regularization parameter; (random) percentage of nodes dropped out 
            of the training.

        direct_link: boolean.
            indicates if the original predictors are included (True) in model's fitting or not (False).

        n_clusters: int.
            number of clusters for 'kmeans' or 'gmm' clustering (could be 0: no clustering).

        cluster_encode: bool.
            defines how the variable containing clusters is treated (default is one-hot)
            if `False`, then labels are used, without one-hot encoding.

        type_clust: str.
            type of clustering method: currently k-means ('kmeans') or Gaussian 
            Mixture Model ('gmm').

        type_scaling: a tuple of 3 strings.
            scaling methods for inputs, hidden layer, and clustering respectively
            (and when relevant). 
            Currently available: standardization ('std') or MinMax scaling ('minmax').

        col_sample: float.
            percentage of covariates randomly chosen for training. 

        lags: int.
            number of lags used for each time series. 

        alpha1: float.
            smoothing weight (training set).

        alpha2: float.
            smoothing parameter (testing set). 

        seed: int. 
            reproducibility seed for nodes_sim=='uniform'.
            
        backend: str.
            "cpu" or "gpu" or "tpu".                           
    """

    # construct the object -----

    def __init__(
        self,
        obj,
        n_hidden_features=5,
        activation_name="relu",
        a=0.01,
        nodes_sim="sobol",
        bias=True,
        dropout=0,
        direct_link=True,
        n_clusters=2,
        cluster_encode=True,
        type_clust="kmeans",
        type_scaling=("std", "std", "std"),
        col_sample=1,
        lags=1,
        alpha1=None,
        alpha2=None,
        seed=123,
        backend="cpu"
    ):

        assert np.int(lags) == lags, "parameter 'lags' should be an integer"

        super().__init__(
            n_hidden_features=n_hidden_features,
            activation_name=activation_name,
            a=a,
            nodes_sim=nodes_sim,
            bias=bias,
            dropout=dropout,
            direct_link=direct_link,
            n_clusters=n_clusters,
            cluster_encode=cluster_encode,
            type_clust=type_clust,
            type_scaling=type_scaling,
            col_sample=col_sample,
            seed=seed,
            backend=backend
        )

        self.obj = obj
        self.n_series = None
        self.lags = lags
        self.fit_objs = {}
        self.y = None  # MTS responses (most recent observations first)
        self.X = None  # MTS lags
        self.xreg = None
        self.y_means = {}
        self.preds = None
        self.preds_std = []
        self.row_sample = 1
        self.alpha1 = alpha1
        self.alpha2 = alpha2

    def fit(self, X, xreg=None, **kwargs):
        """Fit MTS model to training data X, with optional regressors xreg
        
        Args:
        
            X: {array-like}, shape = [n_samples, n_features]
                Training time series, where n_samples is the number 
                of samples and n_features is the number of features;
                X must be in increasing order (most recent observations last)

            xreg: {array-like}, shape = [n_samples, n_features_xreg]
                Additional regressors to be passed to obj
                xreg must be in increasing order (most recent observations last)

            **kwargs: additional parameters to be passed to 
                    self.cook_training_set
               
        Returns: 
        
            self: object
        """

        try:
            # multivariate time series
            n, p = X.shape
        except:
            # univariate time series
            n = X.shape[0]
            p = 1

        rep_1_n = np.repeat(1, n)

        self.y = None
        self.X = None
        self.n_series = p
        self.fit_objs.clear()
        self.y_means.clear()

        if p > 1:
            # multivariate time series
            mts_input = ts.create_train_inputs(X[::-1], self.lags)
        else:
            # univariate time series
            mts_input = ts.create_train_inputs(
                X.reshape(-1, 1)[::-1], self.lags
            )

        self.y = mts_input[0]
        # print(f"self.y: \n {self.y} \n ")
        # print(f"mts_input[1]: \n {mts_input[1]} \n ")
        if self.alpha1 is not None:
            alphas_seq = np.asarray(
                [
                    self.alpha1 * ((1 - self.alpha1) ** i)
                    for i in range(n - 1, -1, -1)
                ]
            )  # first to last
            alphas_seq_lags = ts.create_lags(
                alphas_seq, self.lags - 1, self.n_series
            )[:-1, :]
            # print(f"alphas_seq_lags: \n {alphas_seq_lags}")
            self.X = mts_input[1] * np.exp(alphas_seq_lags)
        else:
            self.X = mts_input[1]

        if xreg is not None:

            assert (
                xreg.shape[0] == n
            ), "'xreg' and 'X' must have the same number or observations"

            self.xreg = xreg

            xreg_input = ts.create_train_inputs(xreg[::-1], self.lags)

            dummy_y, scaled_Z = self.cook_training_set(
                y=rep_1_n,
                X=mo.cbind(self.X, xreg_input[1], backend=self.backend),
                **kwargs
            )

        else:  # xreg is None

            # avoids scaling X p times in the loop
            dummy_y, scaled_Z = self.cook_training_set(
                y=rep_1_n, X=self.X, **kwargs
            )

        # loop on all the time series and adjust self.obj.fit
        for i in range(p):
            y_mean = np.mean(self.y[:, i])
            self.y_means[i] = y_mean
            self.fit_objs[i] = pickle.loads(
                pickle.dumps(
                    self.obj.fit(scaled_Z, self.y[:, i] - y_mean, **kwargs), -1
                )
            )

        return self

    def predict(self, h=5, level=95, new_xreg=None, **kwargs):
        """Forecast all the time series, h steps ahead
        
        Args:
        
            h: {integer}
                Forecasting horizon
            
            level: {integer}
                Level of confidence (if obj has option 'return_std' and the 
                posterior is gaussian)
                
            new_xreg: {array-like}, shape = [n_samples = h, n_new_xreg]
                New values of additional (deterministic) regressors on horizon = h
                new_xreg must be in increasing order (most recent observations last)
            
            **kwargs: additional parameters to be passed to 
                    self.cook_test_set
               
        Returns: 
        
            model predictions for horizon = h: {array-like}
        """

        self.return_std = False

        if self.xreg is not None:
            assert new_xreg is not None, "'new_xreg' must be provided"

        self.preds = None  # do not remove (!)

        self.preds = pickle.loads(pickle.dumps(self.y, -1))

        if self.alpha2 is not None:
            if len(self.preds) > 1:
                # multivariate
                previous_pred = self.preds[0, :]
            else:
                # univariate
                previous_pred = self.preds[0]

        n_features = self.n_series * self.lags

        if "return_std" in kwargs:

            self.preds_std = np.zeros(h)

            self.return_std = kwargs["return_std"]

            qt = norm.ppf(1 - 0.5 * (1 - level / 100))

        if new_xreg is not None:  # Additional regressors provided

            try:
                n_obs_xreg, n_features_xreg = new_xreg.shape
                assert (
                    n_features_xreg == self.xreg.shape[1]
                ), "check number of series provided for 'new_xreg' (compare with self.xreg.shape[1])"
            except:
                n_obs_xreg = new_xreg.shape  # one series

            assert (
                n_obs_xreg == h
            ), "please provide values of regressors 'new_xreg' for the whole horizon 'h'"

            n_features_xreg = n_features_xreg * self.lags

            n_features_total = n_features + n_features_xreg

            inv_new_xreg = mo.rbind(self.xreg, new_xreg)[::-1]

            # Loop on horizon h
            for i in range(h):

                new_obs = ts.reformat_response(self.preds, self.lags)

                new_obs_xreg = ts.reformat_response(inv_new_xreg, self.lags)

                new_X = mo.rbind(
                    np.union1d(
                        new_obs.reshape(1, n_features),
                        new_obs_xreg.reshape(1, n_features_xreg),
                    ),
                    np.ones(n_features_total).reshape(1, n_features_total),
                )

                cooked_new_X = self.cook_test_set(new_X, **kwargs)

                predicted_cooked_new_X = self.obj.predict(
                    cooked_new_X, **kwargs
                )

                if self.return_std == False:  # std. dev. is not returned

                    if self.alpha2 is not None:

                        preds = (
                            self.alpha2
                            * np.array(
                                [
                                    (
                                        self.y_means[j]
                                        + predicted_cooked_new_X[0]
                                    )
                                    for j in range(self.n_series)
                                ]
                            )
                            + (1 - self.alpha2) * previous_pred
                        )

                    else:

                        preds = np.array(
                            [
                                (self.y_means[j] + predicted_cooked_new_X[0])
                                for j in range(self.n_series)
                            ]
                        )

                else:  # std. dev. is returned

                    if self.alpha2 is not None:

                        preds = (
                            self.alpha2
                            * np.array(
                                [
                                    (
                                        self.y_means[j]
                                        + predicted_cooked_new_X[0][0]
                                    )
                                    for j in range(self.n_series)
                                ]
                            )
                            + (1 - self.alpha2) * previous_pred
                        )

                    else:

                        preds = np.array(
                            [
                                (self.y_means[j] + predicted_cooked_new_X[0][0])
                                for j in range(self.n_series)
                            ]
                        )

                    self.preds_std[i] = predicted_cooked_new_X[1][0]

                if self.alpha2 is not None:

                    previous_pred = pickle.loads(pickle.dumps(preds, -1))

                self.preds = mo.rbind(preds, self.preds)

        else:  # No additional regressor provided

            for i in range(h):

                new_obs = ts.reformat_response(self.preds, self.lags)

                new_X = mo.rbind(
                    new_obs.reshape(1, n_features),
                    np.ones(n_features).reshape(1, n_features),
                )

                cooked_new_X = self.cook_test_set(new_X, **kwargs)

                predicted_cooked_new_X = self.obj.predict(
                    cooked_new_X, **kwargs
                )

                if self.return_std == False:  # std. dev. is not returned

                    if self.alpha2 is not None:

                        preds = (
                            self.alpha2
                            * np.array(
                                [
                                    (
                                        self.y_means[j]
                                        + predicted_cooked_new_X[0]
                                    )
                                    for j in range(self.n_series)
                                ]
                            )
                            + (1 - self.alpha2) * previous_pred
                        )

                    else:

                        preds = np.array(
                            [
                                (self.y_means[j] + predicted_cooked_new_X[0])
                                for j in range(self.n_series)
                            ]
                        )

                else:  # std. dev. is returned

                    if self.alpha2 is not None:

                        preds = (
                            self.alpha2
                            * np.array(
                                [
                                    (
                                        self.y_means[j]
                                        + predicted_cooked_new_X[0][0]
                                    )
                                    for j in range(self.n_series)
                                ]
                            )
                            + (1 - self.alpha2) * previous_pred
                        )

                    else:

                        preds = np.array(
                            [
                                (self.y_means[j] + predicted_cooked_new_X[0][0])
                                for j in range(self.n_series)
                            ]
                        )

                    self.preds_std[i] = predicted_cooked_new_X[1][0]

                if self.alpha2 is not None:
                    previous_pred = pickle.loads(pickle.dumps(preds, -1))

                self.preds = mo.rbind(preds, self.preds)

        # function's return

        self.preds = self.preds[0:h, :][::-1]

        if self.return_std == False:  # std. dev. is returned

            return self.preds

        # std. dev. is not returned
        self.preds_std = self.preds_std[::-1].reshape(h, 1)

        self.preds_std = np.repeat(self.preds_std, self.n_series).reshape(
            -1, self.n_series
        )

        return (
            self.preds,
            self.preds_std,
            self.preds - qt * self.preds_std,
            self.preds + qt * self.preds_std,
        )

    def score(self, X, training_index, testing_index, scoring=None, **kwargs):
        """ Train on training_index, score on testing_index. """

        assert (
            bool(set(training_index).intersection(set(testing_index))) == False
        ), "Non-overlapping 'training_index' and 'testing_index' required"

        # Dimensions
        n, p = X.shape

        # Training and testing sets
        X_train = X[training_index, :]
        X_test = X[testing_index, :]

        # Horizon
        h = len(testing_index)
        assert (
            len(training_index) + h
        ) <= n, "Please check lengths of training and testing windows"

        # Fit and predict
        self.fit(X_train, **kwargs)
        preds = self.predict(h=h, **kwargs)

        if scoring is None:
            scoring = "neg_root_mean_squared_error"

        # check inputs
        assert scoring in (
            "explained_variance",
            "neg_mean_absolute_error",
            "neg_mean_squared_error",
            "neg_root_mean_squared_error",
            "neg_mean_squared_log_error",
            "neg_median_absolute_error",
            "r2",
        ), "'scoring' should be in ('explained_variance', 'neg_mean_absolute_error', \
                               'neg_mean_squared_error', 'neg_root_mean_squared_error', 'neg_mean_squared_log_error', \
                               'neg_median_absolute_error', 'r2')"

        scoring_options = {
            "explained_variance": skm2.explained_variance_score,
            "neg_mean_absolute_error": skm2.mean_absolute_error,
            "neg_mean_squared_error": skm2.mean_squared_error,
            "neg_root_mean_squared_error": lambda x, y: np.sqrt(
                skm2.mean_squared_error(x, y)
            ),
            "neg_mean_squared_log_error": skm2.mean_squared_log_error,
            "neg_median_absolute_error": skm2.median_absolute_error,
            "r2": skm2.r2_score,
        }

        return tuple(
            [
                scoring_options[scoring](X_test[:, i], preds[:, i], **kwargs)
                for i in range(p)
            ]
        )
