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

__all__ = ['ESRNN']

# Cell
from typing import Union, List

import numpy as np
import pytorch_lightning as pl
import torch as t
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR

from .esrnn_model import _ESRNN
from ...losses.utils import LossFunction

# Cell
class ESRNN(pl.LightningModule):
    def __init__(self, n_series: int,
                 n_x: int, n_s: int,
                 input_size: int, output_size: int,
                 idx_to_sample_freq: int,
                 es_component: str = 'multiplicative',
                 cell_type: str = 'LSTM', state_hsize: int = 50,
                 dilations: List[List[int]] = [[1, 2], [4, 8]],
                 add_nl_layer: bool = False, seasonality: List[int] = [],
                 learning_rate: float = 1e-3, lr_scheduler_step_size: int = 9,
                 lr_decay: float = 0.9, per_series_lr_multip: float = 1.,
                 gradient_eps: float = 1e-8,
                 gradient_clipping_threshold: float = 20.,
                 rnn_weight_decay: float = 0., noise_std: float = 1e-3,
                 level_variability_penalty: float = 20.,
                 testing_percentile: Union[int, List] = 50,
                 training_percentile: Union[int, List] = 50,
                 loss: str = 'SMYL', val_loss: str = 'MAE'):
        super(ESRNN, self).__init__()
        """ Exponential Smoothing Recurrent Neural Network

        Pytorch Implementation of the M4 time series forecasting competition winner.
        Proposed by Smyl. The model uses a hybrid approach of Machine Learning and
        statistical methods by combining recurrent neural networks to model a common
        trend with shared parameters across series, and multiplicative Holt-Winter
        exponential smoothing.

        Parameters
        ----------
        n_series: int
            Number of time series.
        n_x: int
            Number of temporal exogenous variables.
        n_s: int
            Number of static variables.
        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
        idx_to_sample_freq: int
            Step size between windows.
        es_component: str
            Exponential Smoothing component.
            Default multiplicative.
        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
        seasonality: int list
            list of seasonalities of the time series
            Hourly [24, 168], Daily [7], Weekly [52], Monthly [12],
            Quarterly [4], Yearly [].
        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.
        per_series_lr_multip: float
            multiplier for per-series parameters smoothing and initial
            seasonalities learning rate (default 1.0)
        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
        rnn_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
        level_variability_penalty: float
            this parameter controls the strength of the penalization
            to the wigglines of the level vector, induces smoothness
            in the output
        testing_percentile: int
            This value is only for diagnostic evaluation.
            In case of percentile predictions this parameter controls
            for the value predicted, when forecasting point value,
            the forecast is the median, so percentile=50.
        training_percentile: float
            To reduce the model's tendency to over estimate, the
            training_percentile can be set to fit a smaller value
            through the Pinball Loss.
        loss: str
            Loss used to train.
        val_loss: str
            Loss used to validate.

        Notes
        -----
        **References:**
        `M4 Competition Conclusions
        <https://rpubs.com/fotpetr/m4competition>`__
        `Original Dynet Implementation of ESRNN
        <https://github.com/M4Competition/M4-methods/tree/master/118%20-%20slaweks17>`__
        """

        #------------------------ Model Attributes ------------------------#
        # Architecture parameters
        self.n_series = n_series
        self.n_x = n_x
        self.n_s = n_s
        self.idx_to_sample_freq = idx_to_sample_freq
        self.input_size = input_size
        self.output_size = output_size
        self.es_component = es_component
        self.cell_type = cell_type
        self.state_hsize = state_hsize
        self.dilations = dilations
        self.add_nl_layer = add_nl_layer
        self.seasonality = seasonality

        # Regularization and optimization parameters
        self.learning_rate = learning_rate
        self.lr_scheduler_step_size = lr_scheduler_step_size
        self.lr_decay = lr_decay
        self.per_series_lr_multip = per_series_lr_multip
        self.gradient_eps = gradient_eps
        self.gradient_clipping_threshold = gradient_clipping_threshold
        self.rnn_weight_decay = rnn_weight_decay
        self.noise_std = noise_std
        self.level_variability_penalty = level_variability_penalty
        self.testing_percentile = testing_percentile
        self.training_percentile = training_percentile
        self.loss = loss
        self.val_loss = val_loss
        self.loss_fn = LossFunction(loss,
                                    percentile=self.training_percentile,
                                    level_variability_penalty=self.level_variability_penalty)
        self.val_loss_fn = LossFunction(val_loss,
                                        percentile=self.testing_percentile,
                                        level_variability_penalty=self.level_variability_penalty)

        # MQESRNN
        self.mq = isinstance(self.training_percentile, list)
        self.output_size_m = len(self.training_percentile) if self.mq else 1

        #Defining model
        self.esrnn = _ESRNN(n_series=self.n_series,
                            input_size=self.input_size,
                            output_size=self.output_size,
                            output_size_m=self.output_size_m,
                            n_t=self.n_x, n_s=self.n_s,
                            es_component=self.es_component,
                            seasonality=self.seasonality,
                            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 training_step(self, batch, batch_idx):
        es_opt, rnn_opt = self.optimizers()

        #Parsing batch
        S = batch['S']
        Y = batch['Y']
        X = batch['X']
        idxs = batch['idxs']

        target, forecast, levels = self.esrnn(S=S, Y=Y, X=X, idxs=idxs,
                                              step_size=self.idx_to_sample_freq)
        loss = self.loss_fn(y=target,
                            y_hat=forecast,
                            y_insample=Y,
                            levels=levels)

        es_opt.zero_grad()
        rnn_opt.zero_grad()
        self.manual_backward(loss)
        clip_grad_norm_(parameters=self.esrnn.rnn.parameters(),
                        max_norm=self.gradient_clipping_threshold)
        clip_grad_norm_(parameters=self.esrnn.es.parameters(),
                        max_norm=self.gradient_clipping_threshold)
        es_opt.step()
        rnn_opt.step()

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

        return loss

    def on_train_epoch_end(self):
        for lr_scheduler in self.lr_schedulers():
            lr_scheduler.step()

    def validation_step(self, batch, idx):
        S = batch['S']
        Y = batch['Y']
        X = batch['X']
        idxs = batch['idxs']

        # TODO: improve temporal hack, mask is not used in loss yet
        target, forecast, _ = self.esrnn.predict(S=S, Y=Y, X=X, idxs=idxs,
                                                 step_size=self.idx_to_sample_freq)
        loss = self.val_loss_fn(y=target,
                                y_hat=forecast,
                                y_insample=Y)
        self.log('val_loss', loss,  prog_bar=True)

        return loss

    def forward(self, batch):
        S = batch['S']
        Y = batch['Y']
        X = batch['X']
        idxs = batch['idxs']

        target, forecast, outsample_mask = self.esrnn.predict(S=S, Y=Y, X=X, idxs=idxs,
                                                              step_size=self.idx_to_sample_freq)

        return target, forecast, outsample_mask

    def configure_optimizers(self):
        es_optimizer = Adam(params=self.esrnn.es.parameters(),
                            lr=self.learning_rate * self.per_series_lr_multip,
                            eps=self.gradient_eps)
        rnn_optimizer = Adam(params=self.esrnn.rnn.parameters(),
                             lr=self.learning_rate,
                             eps=self.gradient_eps,
                             weight_decay=self.rnn_weight_decay)

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

        return [es_optimizer, rnn_optimizer], [lr_es, lr_rnn]