# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/models.gmm_tft.ipynb.

# %% auto 0
__all__ = ['GMM_TFT']

# %% ../../nbs/models.gmm_tft.ipynb 2
import torch
import torch.nn as nn

import logging
import warnings
logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")

from .tft import TFT
from ..losses.pytorch import MAE

# %% ../../nbs/models.gmm_tft.ipynb 4
class GMM_TFT(TFT):

    def __init__(self,
                 input_size,
                 h,
                 K=100,
                 tgt_size=1,
                 hidden_size=128,
                 s_cont_cols=None,
                 s_cat_cols=None,
                 o_cont_cols=None,
                 o_cat_cols=None,
                 k_cont_cols=None,
                 k_cat_cols=None,
                 s_cat_inp_lens=None,
                 s_cont_inp_size=0,
                 k_cat_inp_lens=None,
                 k_cont_inp_size=1,
                 o_cat_inp_lens=None,
                 o_cont_inp_size=0,
                 n_head=4,
                 attn_dropout=0.0,
                 dropout=0.1,
                 windows_batch_size=1024,
                 step_size=1,
                 learning_rate=1e-3,
                 normalize=True,
                 loss=MAE(),
                 batch_size=32, 
                 num_workers_loader=0,
                 drop_last_loader=False,
                 random_seed=1,
                 **trainer_kwargs):
        # Inherit TFT class and extend it with GMM
        super(GMM_TFT, self).__init__(input_size=input_size,
                                      h=h,
                                      tgt_size=tgt_size,
                                      hidden_size=hidden_size,
                                      s_cont_cols=s_cont_cols,
                                      s_cat_cols=s_cat_cols,
                                      o_cont_cols=o_cont_cols,
                                      o_cat_cols=o_cat_cols,
                                      k_cont_cols=k_cont_cols,
                                      k_cat_cols=k_cat_cols,
                                      s_cat_inp_lens=s_cat_inp_lens,
                                      s_cont_inp_size=s_cont_inp_size,
                                      k_cat_inp_lens=k_cat_inp_lens,
                                      k_cont_inp_size=k_cont_inp_size,
                                      o_cat_inp_lens=o_cat_inp_lens,
                                      o_cont_inp_size=o_cont_inp_size,
                                      n_head=n_head,
                                      attn_dropout=attn_dropout,
                                      dropout=dropout,
                                      windows_batch_size=windows_batch_size,
                                      step_size=step_size,
                                      learning_rate=learning_rate,
                                      normalize=normalize,
                                      loss=loss,
                                      batch_size=batch_size ,
                                      num_workers_loader=num_workers_loader,
                                      drop_last_loader=drop_last_loader,
                                      random_seed=random_seed,
                                      **trainer_kwargs)
        
        # Define Mixture specialized parameters
        self.K = K

        # Adapter with Loss dependent dimensions
        self.output_adapter = nn.Linear(in_features=hidden_size, out_features=K)
    
    def training_step(self, batch, batch_idx):
        # Deviates from orignal `BaseWindows.training_step` to 
        # allow the model to receive future exogenous available
        # at the time of the prediction.
        
        # Create windows [Ws, L+H, C]
        windows = self._create_windows(batch, step='train')

        # Normalize
        if self.normalize:
            windows, *_ = self._normalization(windows)

        # outsample
        y_idx = batch['temporal_cols'].get_loc('y')
        mask_idx = batch['temporal_cols'].get_loc('available_mask')
        outsample_y = windows['temporal'][:, -self.h:, y_idx]
        outsample_mask = windows['temporal'][:, -self.h:, mask_idx]

        # [Ws, H, K]
        means_hat = self(x=windows)
        stds = (1/self.K) * torch.ones_like(means_hat).to(means_hat.device)
        weights = (1/self.K) * torch.ones_like(means_hat).to(means_hat.device)

        loss = self.loss(y=outsample_y[:,:,None], means=means_hat,
                         stds=stds, weights=weights,
                         mask=outsample_mask[:,:,None])
        self.log('train_loss', loss, prog_bar=True, on_epoch=True)
        return loss

    def predict_step(self, batch, batch_idx):
        # Deviates from orignal `BaseWindows.training_step` to 
        # allow the model to receive future exogenous available
        # at the time of the prediction.        
        
        # Create windows [Ws, L+H, C]
        windows = self._create_windows(batch, step='predict')

        # Normalize windows
        if self.normalize:
            windows, y_means, y_stds = self._normalization(windows)

        # [Ws, H, K]
        means_hat = self(x=windows)
        stds = (1/self.K) * torch.ones_like(means_hat).to(means_hat.device)
        weights = (1/self.K) * torch.ones_like(means_hat).to(means_hat.device)

        _, quants = self.loss.sample(weights=weights,
                                     means=means_hat, stds=stds,
                                     num_samples=2000)

        # Inv Normalize [Ws, H, Q]
        if self.normalize:
            quants = self._inv_normalization(quants, y_means, y_stds)

        return quants
