# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/models_rnn__rnn.ipynb (unless otherwise specified).

__all__ = ['RNN', 'suggested_space']

# Cell
import numpy as np
import random
import torch as t
import torch.nn as nn
from hyperopt import hp

from typing import Union, List
from ..components.drnn import DRNN

# Cell
class _RNN(nn.Module):
    def __init__(self, input_size: int, output_size: int,
                 n_t: int, n_s: int, cell_type: str, dilations: list, state_hsize: int, add_nl_layer: bool):
        super(_RNN, self).__init__()

        self.input_size = input_size
        self.output_size = output_size
        self.n_t = n_t
        self.n_s = n_s
        self.cell_type = cell_type
        self.dilations = dilations
        self.state_hsize = state_hsize
        self.add_nl_layer = add_nl_layer
        self.layers = len(dilations)

        layers = []
        for grp_num in range(len(self.dilations)):
            if grp_num == 0:
                input_size = self.input_size + (self.input_size + self.output_size)*self.n_t + self.n_s
            else:
                input_size = self.state_hsize
            layer = DRNN(input_size,
                         self.state_hsize,
                         n_layers=len(self.dilations[grp_num]),
                         dilations=self.dilations[grp_num],
                         cell_type=self.cell_type)
            layers.append(layer)

        self.rnn_stack = nn.Sequential(*layers)

        if self.add_nl_layer:
            self.MLPW  = nn.Linear(self.state_hsize, self.state_hsize)

        self.adapterW  = nn.Linear(self.state_hsize, self.output_size)

    def forward(self, Y: t.Tensor, X: t.Tensor):
        if self.n_t >0:
            input_data = t.cat((Y, X), -1)
        else:
            input_data = Y

        for layer_num in range(len(self.rnn_stack)):
            residual = input_data
            output, _ = self.rnn_stack[layer_num](input_data)
            if layer_num > 0:
                output += residual
            input_data = output

        if self.add_nl_layer:
            input_data = self.MLPW(input_data)
            input_data = t.tanh(input_data)

        input_data = self.adapterW(input_data)
        input_data = input_data.transpose(0,1) #change to bs, n_windows

        return input_data

# Cell
from typing import Union, List

import pandas as pd
import pytorch_lightning as pl
import torch.nn as nn
from fastcore.foundation import patch
from torch.nn.utils import clip_grad_norm_
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR

from ...data.tsdataset import TimeSeriesDataset
from ...data.tsloader import TimeSeriesLoader
from ...losses.utils import LossFunction

# Cell
class RNN(pl.LightningModule):
    def __init__(self,
                 input_size: int, output_size: int,
                 n_x: int = 0, n_s: int = 0,
                 sample_freq: int = 1,
                 cell_type: str = 'LSTM', state_hsize: int = 50,
                 dilations: List[List[int]] = [[1, 2], [4, 8]],
                 add_nl_layer: bool = False,
                 learning_rate: float = 1e-3, lr_scheduler_step_size: int = 1000,
                 lr_decay: float = 0.9,
                 gradient_eps: float = 1e-8,
                 gradient_clipping_threshold: float = 20.,
                 weight_decay: float = 0., noise_std: float = 1e-3,
                 loss_train: str = 'MAE', loss_valid: str = 'MAE',
                 loss_hypar: float = 0.,
                 frequency: str = 'D',
                 random_seed: int = 1):
        super(RNN, self).__init__()

        """ Recurrent Neural Network

        Parameters
        ----------
        input_size: int
            input size of the recurrent neural network, usually a
            multiple of seasonality
        output_size: int
            output_size or forecast horizon of the recurrent neural
            network, usually multiple of seasonality
        n_x: int
            Number of temporal exogenous variables.
        n_s: int
            Number of static variables.
        sample_freq: int
            Step size between windows.
        cell_type: str
            Type of RNN cell, available GRU, LSTM, RNN, ResidualLSTM.
        state_hsize: int
            dimension of hidden state of the recurrent neural network
        dilations: int list
            each list represents one chunk of Dilated LSTMS, connected in
            standard ResNet fashion
        add_nl_layer: bool
            whether to insert a tanh() layer between the RNN stack and the
            linear adaptor (output) layers
        learning_rate: float
            size of the stochastic gradient descent steps
        lr_scheduler_step_size: int
            this step_size is the period for each learning rate decay
        lr_decay: float
            Learning rate decay.
        gradient_eps: float
            term added to the Adam optimizer denominator to improve
            numerical stability (default: 1e-8)
        gradient_clipping_threshold: float
            max norm of gradient vector, with all parameters treated
            as a single vector
        weight_decay: float
            parameter to control classic L2/Tikhonov regularization
            of the rnn parameters
        noise_std: float
            standard deviation of white noise added to input during
            fit to avoid the model from memorizing the train data
        loss_train: str
            Loss used to train.
        loss_valid: str
            Loss used to validate.
        loss_hypar: float
            Hyperparameter for chosen loss.
        frequency: str
            Time series frequency.
        random_seed: int
            random_seed for pseudo random pytorch initializer and
            numpy random generator.

        """

        #------------------------ Model Attributes ------------------------#
        # Architecture parameters
        self.n_x = n_x
        self.n_s = n_s
        self.sample_freq = sample_freq
        self.input_size = input_size
        self.output_size = output_size
        self.cell_type = cell_type
        self.state_hsize = state_hsize
        self.dilations = dilations
        self.add_nl_layer = add_nl_layer

        # Regularization and optimization parameters
        self.learning_rate = learning_rate
        self.lr_scheduler_step_size = lr_scheduler_step_size
        self.lr_decay = lr_decay
        self.gradient_eps = gradient_eps
        self.gradient_clipping_threshold = gradient_clipping_threshold
        self.weight_decay = weight_decay
        self.noise_std = noise_std
        self.random_seed = random_seed

        # Loss functions
        self.loss_train = loss_train
        self.loss_hypar = loss_hypar
        self.loss_valid = loss_valid
        self.loss_fn_train = LossFunction(loss_train,
                                          percentile=self.loss_hypar,
                                          seasonality=self.loss_hypar)
        self.loss_fn_valid = LossFunction(loss_valid,
                                          percentile=self.loss_hypar,
                                          seasonality=self.loss_hypar)

        self.frequency = frequency

        #Defining model
        self.model = _RNN(input_size=self.input_size,
                          output_size=self.output_size,
                          n_t=self.n_x, n_s=self.n_s,
                          #noise_std=self.noise_std,
                          cell_type=self.cell_type,
                          dilations=self.dilations,
                          state_hsize=self.state_hsize,
                          add_nl_layer=self.add_nl_layer)

        self.automatic_optimization = False

    def on_fit_start(self):
        t.manual_seed(self.random_seed)
        np.random.seed(self.random_seed)
        random.seed(self.random_seed)

    def parse_batch(self, batch):
        S = batch['S']
        Y = batch['Y']
        X = batch['X']
        idxs = batch['idxs']
        sample_mask = batch['sample_mask']
        available_mask = batch['available_mask']

        av_condition = t.nonzero(t.min(available_mask, axis=0).values)
        min_time_stamp = int(av_condition.min())
        sample_condition = t.nonzero(t.min(sample_mask, axis=0).values)

        if sample_condition.nelement() == 0:
            max_time_stamp = int(av_condition.max())
        else:
            max_time_stamp = int(sample_condition.max())
        available_ts = max_time_stamp - min_time_stamp + 1 # +1, inclusive counting
        if available_ts < self.input_size + self.output_size:
            raise Exception(
                'Time series too short for given input and output size. \n'
                f'Available timestamps: {available_ts}'
            )

        Y = Y[:, min_time_stamp:max_time_stamp + 1] #+1 because is not inclusive
        X = X[:, :, min_time_stamp:max_time_stamp + 1]
        sample_mask = sample_mask[:, min_time_stamp:max_time_stamp + 1]
        available_mask = available_mask[:, min_time_stamp:max_time_stamp + 1]

        Y = Y.unfold(dimension=-1, size=self.input_size+self.output_size, step=self.sample_freq)
        X = X.unfold(dimension=-1, size=self.input_size+self.output_size, step=self.sample_freq)
        sample_mask = sample_mask.unfold(dimension=-1, size=self.input_size+self.output_size, step=self.sample_freq)
        available_mask = available_mask.unfold(dimension=-1, size=self.input_size+self.output_size, step=self.sample_freq)

        Y = Y.transpose(0,1) # n_windows, batch, time
        X = X.transpose(1,2) # batch, n_windows, n_s, time
        X = X.flatten(start_dim=2)
        X = X.transpose(0,1) # n_windows, batch, n_s*time
        sample_mask = sample_mask.transpose(0,1)
        available_mask = available_mask.transpose(0,1)

        return S, Y, X, idxs, sample_mask, available_mask

    def training_step(self, batch, batch_idx):
        rnn_opt = self.optimizers()

        # Parsing batch
        S, Y, X, idxs, sample_mask, available_mask = self.parse_batch(batch)

        insample_Y = Y[:,:, :self.input_size]
        outsample_Y = Y[:,:, self.input_size:]
        outsample_mask = sample_mask[:,:, self.input_size:]
        outsample_Y = outsample_Y.transpose(0,1)
        outsample_mask = outsample_mask.transpose(0,1)

        y_hat = self.model(Y=insample_Y, X=X)

        loss = self.loss_fn_train(y=outsample_Y,
                                  y_hat=y_hat,
                                  y_insample=Y,
                                  mask=outsample_mask)

        rnn_opt.zero_grad()
        self.manual_backward(loss)
        clip_grad_norm_(parameters=self.model.parameters(),
                        max_norm=self.gradient_clipping_threshold)
        rnn_opt.step()

        self.log('train_loss', loss, prog_bar=True, on_epoch=True)

        return loss

    def on_train_epoch_end(self):
        self.lr_schedulers().step()

    def validation_step(self, batch, idx):
        # Parsing batch
        S, Y, X, idxs, sample_mask, available_mask = self.parse_batch(batch)

        insample_Y = Y[:,:, :self.input_size]
        outsample_Y = Y[:,:, self.input_size:]
        outsample_mask = sample_mask[:,:, self.input_size:]
        outsample_Y = outsample_Y.transpose(0,1)
        outsample_mask = outsample_mask.transpose(0,1)

        y_hat = self.model(Y=insample_Y, X=X)

        loss = self.loss_fn_valid(y=outsample_Y,
                                  y_hat=y_hat,
                                  mask=outsample_mask,
                                  y_insample=Y)
        self.log('val_loss', loss, prog_bar=True)

        return loss

    def forward(self, batch):
        # Parsing batch
        S, Y, X, idxs, sample_mask, available_mask = self.parse_batch(batch)

        insample_Y = Y[:,:, :self.input_size]
        y_true = Y[:,:, self.input_size:]
        sample_mask = sample_mask[:,:, self.input_size:]
        y_true = y_true.transpose(0,1)
        sample_mask = sample_mask.transpose(0,1)

        y_hat = self.model(Y=insample_Y, X=X)

        # Filter to windows with at least one sampleable ts
        sample_condition = sample_mask.sum(dim=2)>0
        assert t.equal(sample_condition[[0]]*t.ones(sample_condition.shape), 1.0*sample_condition), 'Sample mask must match perfectly between time-series'
        sample_index = sample_condition[0]

        y_true=y_true[:,sample_index,:]
        y_hat=y_hat[:,sample_index,:]
        sample_mask=sample_mask[:,sample_index,:]

        return y_true, y_hat, sample_mask

    def configure_optimizers(self):
        rnn_optimizer = Adam(params=self.model.parameters(),
                             lr=self.learning_rate,
                             eps=self.gradient_eps,
                             weight_decay=self.weight_decay)

        lr_rnn = StepLR(optimizer=rnn_optimizer,
                        step_size=self.lr_scheduler_step_size,
                        gamma=self.lr_decay)

        return [rnn_optimizer], [lr_rnn]

# Cell
@patch
def forecast(self: RNN, Y_df, X_df = None, S_df = None, batch_size=1, trainer=None):
    """
    Method for forecasting self.output_size periods after last timestamp of Y_df.

    Parameters
    ----------
    Y_df: pd.DataFrame
        Dataframe with target time-series data, needs 'unique_id','ds' and 'y' columns.
    X_df: pd.DataFrame
        Dataframe with exogenous time-series data, needs 'unique_id' and 'ds' columns.
        Note that 'unique_id' and 'ds' must match Y_df plus the forecasting horizon.
    S_df: pd.DataFrame
        Dataframe with static data, needs 'unique_id' column.
    bath_size: int
        Batch size for forecasting.

    Returns
    ----------
    forecast_df: pd.DataFrame
        Dataframe with forecasts.
    """

    # Add forecast dates to Y_df
    Y_df['ds'] = pd.to_datetime(Y_df['ds'])
    if X_df is not None:
        X_df['ds'] = pd.to_datetime(X_df['ds'])
    forecast_dates = pd.date_range(Y_df['ds'].max(), periods=self.output_size+1, freq=self.frequency)[1:]
    index = pd.MultiIndex.from_product([Y_df['unique_id'].unique(), forecast_dates], names=['unique_id', 'ds'])
    forecast_df = pd.DataFrame({'y':[0]}, index=index).reset_index()

    Y_df = Y_df.append(forecast_df).sort_values(['unique_id','ds']).reset_index(drop=True)

    # Dataset, loader and trainer
    dataset = TimeSeriesDataset(Y_df=Y_df, X_df=X_df,
                                S_df=S_df,
                                ds_in_test=self.output_size,
                                is_test=True,
                                input_size=self.input_size,
                                output_size=self.output_size,
                                verbose=True)

    loader = TimeSeriesLoader(dataset=dataset,
                              batch_size=batch_size,
                              shuffle=False)

    if trainer is None:
        gpus = -1 if t.cuda.is_available() else 0
        trainer = pl.Trainer(progress_bar_refresh_rate=1,
                             gpus=gpus,
                             logger=False)

    # Forecast
    outputs = trainer.predict(self, loader)

    # Process forecast and include in forecast_df
    _, forecast, _ = zip(*outputs)
    forecast = t.cat([forecast_[:, -1] for forecast_ in forecast]).cpu().numpy()
    forecast_df['y'] = forecast.flatten()

    return forecast_df


# Cell
def suggested_space(n_time_out: int, n_series: int, n_x: int, n_s: int,
                    frequency: str) -> dict:
    """
    Suggested hyperparameters search space for tuning. To be used with hyperopt library.

    Parameters
    ----------
    n_time_out: int
        Forecasting horizon.
    n_series: int
        Number of time-series.
    n_x: int
        Number of exogenous variables.
    n_s: int
        Number of static variables.
    frequency: str
        Frequency of time-seris.

    Returns
    ----------
    space: Dict
        Dictionary with search space for hyperopt library.
    """

    space= {# Architecture parameters
        'model':'rnn',
        'mode': 'full',
        'n_time_in': hp.choice('n_time_in', [2*n_time_out, 3*n_time_out]),
        'n_time_out': hp.choice('n_time_out', [n_time_out]),
        'n_x': hp.choice('n_x', [n_x]),
        'n_s': hp.choice('n_s', [n_s]),
        'cell_type': hp.choice('cell_type', ['LSTM', 'GRU']),
        'state_hsize': hp.choice('state_hsize', [10, 20, 50, 100]),
        'dilations': hp.choice('dilations', [ [[1, 2]], [[1, 2, 4, 8]], [[1,2],[4,8]] ]),
        'add_nl_layer': hp.choice('add_nl_layer', [ False ]),
        'sample_freq': hp.choice('sample_freq', [n_time_out]),
        # Regularization and optimization parameters
        'learning_rate': hp.choice('learning_rate', [0.0001, 0.001, 0.005, 0.01, 0.05, 0.1]),
        'lr_decay': hp.choice('lr_decay', [0.5] ),
        'n_lr_decays': hp.choice('n_lr_decays', [3]),
        'gradient_eps': hp.choice('gradient_eps', [1e-8]),
        'gradient_clipping_threshold': hp.choice('gradient_clipping_threshold', [10]),
        'weight_decay': hp.choice('weight_decay', [0]),
        'noise_std': hp.choice('noise_std', [0.001]),
        'max_epochs': hp.choice('max_epochs', [None]),
        'max_steps': hp.choice('max_steps', [500, 1000]),
        'early_stop_patience': hp.choice('early_stop_patience', [10]),
        'eval_freq': hp.choice('eval_freq', [50]),
        'loss_train': hp.choice('loss', ['MAE']),
        'loss_hypar': hp.choice('loss_hypar', [0.5]),
        'loss_valid': hp.choice('loss_valid', ['MAE']),
        # Data parameters
        'normalizer_y': hp.choice('normalizer_y', [None]),
        'normalizer_x': hp.choice('normalizer_x', [None]),
        'complete_windows':  hp.choice('complete_windows', [True]),
        'frequency': hp.choice('frequency', [frequency]),
        'idx_to_sample_freq': hp.choice('idx_to_sample_freq', [1]),
        'val_idx_to_sample_freq': hp.choice('val_idx_to_sample_freq', [1]),
        'batch_size': hp.choice('batch_size', [16, 32, 64]),
        'n_windows': hp.choice('n_windows', [None]),
        'random_seed': hp.quniform('random_seed', 1, 20, 1)}

    return space