# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/common.base_recurrent.ipynb.

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

# %% ../../nbs/common.base_recurrent.ipynb 4
import random

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import TQDMProgressBar

from ..tsdataset import TimeSeriesDataModule

# %% ../../nbs/common.base_recurrent.ipynb 5
class BaseRecurrent(pl.LightningModule):
    
    def __init__(self, 
                 loss,
                 batch_size=32, 
                 num_workers_loader=0,
                 drop_last_loader=False,
                 random_seed=1, 
                 **trainer_kwargs):
        super(BaseRecurrent, self).__init__()
        self.random_seed = random_seed
        pl.seed_everything(self.random_seed, workers=True)
        #fit arguments
        self.val_size: int = 0
        self.test_size: int = 0
        #predict arguments
        self.step_size: int = 1
        
        # Trainer
        # we need to instantiate the trainer each time we want to use it
        self.trainer_kwargs = {**trainer_kwargs}
        if self.trainer_kwargs.get('callbacks', None) is None:
            self.trainer_kwargs = {**{'callbacks': [TQDMProgressBar()], **trainer_kwargs}}
        else:
            self.trainer_kwargs = trainer_kwargs

        # Add GPU accelerator if available
        if self.trainer_kwargs.get('accelerator', None) is None:
            if torch.cuda.is_available():
                self.trainer_kwargs['accelerator'] = "gpu"
        if self.trainer_kwargs.get('devices', None) is None:
            if torch.cuda.is_available():
                self.trainer_kwargs['devices'] = -1
        
        # DataModule arguments
        self.batch_size = batch_size
        self.num_workers_loader = num_workers_loader
        self.drop_last_loader = drop_last_loader
    
    def on_fit_start(self):
        torch.manual_seed(self.random_seed)
        np.random.seed(self.random_seed)
        random.seed(self.random_seed)
        
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)

    def _normalization(self, batch):
        temporal = batch['temporal']
        temporal_cols = batch['temporal_cols']

        # Remove validation and test set to prevent leakeage
        if self.val_size + self.test_size > 0:
                cutoff = -self.val_size - self.test_size
                temporal = temporal[:, :, :cutoff]

        temporal_y = temporal[:, temporal_cols.get_loc('y'), :]
        temporal_mask = temporal[:, temporal_cols.get_loc('available_mask'), :]

        y_means = torch.sum(temporal_y * temporal_mask, dim=-1, keepdim=True) / torch.sum(temporal_mask, dim=-1, keepdim=True)
        y_stds = torch.sqrt(torch.sum(temporal_mask*(temporal_y-y_means)**2, dim=-1, keepdim=True)/ \
                            torch.sum(temporal_mask, dim=-1, keepdim=True) )

        temporal_y = (temporal_y - y_means) / y_stds
        temporal[:, temporal_cols.get_loc('y'), :] = temporal_y
        
        return batch, y_means, y_stds

    def _inv_normalization(self, y_hat, y_means, y_stds):
        return y_stds[:, None]*y_hat + y_means[:, None]

    def _create_windows(self, batch, step):
        temporal = batch['temporal']
        temporal_cols = batch['temporal_cols']

        assert temporal.shape[-1] % self.step_size == 0, f'Temporal length of batch should be multiple of {self.step_size}.'
        assert self.val_size % self.step_size == 0, f'Validation size should be multiple of {self.step_size}.'
        assert self.test_size % self.step_size == 0, f'Test size should be multiple of {self.step_size}.'

        if step == 'train':
            if self.val_size + self.test_size > 0:
                cutoff = -self.val_size - self.test_size
                temporal = temporal[:, :, :cutoff]
            temporal = self.padder(temporal)

            # Truncate batch to shorter time-series 
            av_condition = torch.nonzero(torch.min(temporal[:, temporal_cols.get_loc('available_mask')], axis=0).values)
            min_time_stamp = int(av_condition.min())
            
            available_ts = temporal.shape[-1] - min_time_stamp + 1 # +1, inclusive counting
            if available_ts < self.input_size + self.h:
                raise Exception(
                    'Time series too short for given input and output size. \n'
                    f'Available timestamps: {available_ts}'
                )

            temporal = temporal[:, :, min_time_stamp:]

        if step == 'val':
            if self.test_size > 0:
                temporal = temporal[:, :, :-self.test_size]
            temporal = self.padder(temporal)

        if step == 'predict':
            if self.test_size == 0:
                temporal = self.padder(temporal)

        # Parse batch
        window_size = self.input_size + self.h
        windows = temporal.unfold(dimension=-1,
                                  size=window_size,
                                  step=self.step_size)

        # # [B, C, Ws, L+H] - > Ws, C, B, L+H]
        # windows = windows.transpose(0,2)

        # think about interaction available * sample mask
        windows_batch = dict(Y=windows[:, temporal_cols.get_loc('y'), :, :],
                            available_mask=windows[:, temporal_cols.get_loc('available_mask'), :, :])

        return windows_batch

    def training_step(self, batch, batch_idx):
        # Normalize
        if self.normalize:
            batch, *_ = self._normalization(batch)

        windows = self._create_windows(batch, step='train')

        # insample
        insample_y = windows['Y'][:, :, :-self.h]
        insample_mask = windows['available_mask'][:, :, :-self.h]

        # outsample
        outsample_y = windows['Y'][:, :, -self.h:]
        outsample_mask = windows['available_mask'][:, :, -self.h:]

        y_hat = self(insample_y, insample_mask)

        loss = self.loss(y=outsample_y, y_hat=y_hat, mask=outsample_mask)
        self.log('train_loss', loss, prog_bar=True, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        if self.val_size == 0:
            return np.nan

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

        windows = self._create_windows(batch, step='val')

        # insample
        insample_y = windows['Y'][:, :, :-self.h]
        insample_mask = windows['available_mask'][:, :, :-self.h]
        
        # outsample
        outsample_y = windows['Y'][:, :, -self.h:]
        outsample_mask = windows['available_mask'][:, :, -self.h:]
        
        #batch_size, input_size
        y_hat = self(insample_y, insample_mask)

        # Inv Normalize
        if self.normalize:
            y_hat = self._inv_normalization(y_hat, y_means, y_stds)

        # Remove train y_hat
        val_windows = (self.val_size // self.step_size) + 1
        y_hat = y_hat[:, -val_windows:-1, :]
        outsample_y = outsample_y[:, -val_windows:-1, :]
        outsample_mask = outsample_mask[:, -val_windows:-1, :]

        loss = self.loss(y=outsample_y, y_hat=y_hat, mask=outsample_mask)
        self.log('val_loss', loss, prog_bar=True, on_epoch=True)

        return loss

    def validation_epoch_end(self, outputs):
        if self.val_size == 0:
            return
        avg_loss = torch.stack(outputs).mean()
        self.log("ptl/val_loss", avg_loss)

    def predict_step(self, batch, batch_idx):
        # Normalize
        if self.normalize:
            batch, y_means, y_stds = self._normalization(batch)

        windows = self._create_windows(batch, step='predict')

        # insample
        insample_y = windows['Y'][:, :, :-self.h]
        insample_mask = windows['available_mask'][:, :, :-self.h]

        y_hat = self(insample_y, insample_mask)

        # Inv Normalize
        if self.normalize:
            y_hat = self._inv_normalization(y_hat, y_means, y_stds)

        return y_hat

    def fit(self, dataset, val_size=0, test_size=0):
        """
        Fits Model.
        
        **Parameters:**<br>
        `dataset`: TimeSeriesDataset.<br>
        `trainer`: pl.Trainer.<br>
        `val_size`: int, validation size.<br>
        `test_size`: int, test size.<br>
        `data_kwargs`: extra arguments to be passed to TimeSeriesDataModule.
        """
        self.val_size = val_size
        self.test_size = test_size
        datamodule = TimeSeriesDataModule(
            dataset, 
            batch_size=self.batch_size,
            num_workers=self.num_workers_loader,
            drop_last=self.drop_last_loader
        )
        trainer = pl.Trainer(**self.trainer_kwargs)
        trainer.fit(self, datamodule=datamodule)

    def predict(self, dataset, step_size=1, **data_kwargs):
        """
        Predicts RNN Model.
        
        **Parameters:**<br>
        `dataset`: TimeSeriesDataset.<br>
        `trainer`: pl.Trainer.<br>
        `data_kwargs`: extra arguments to be passed to TimeSeriesDataModule.
        """
        if step_size != self.step_size:
            print(f'WARNING: step_size {step_size} different than self.step_size, \
                   using self.step_size from training instead!')

        self.predict_step_size = self.step_size
        # fcsts (window, batch, h)
        trainer = pl.Trainer(**self.trainer_kwargs)
        fcsts = trainer.predict(self, datamodule=TimeSeriesDataModule(dataset, **data_kwargs))
        if self.test_size > 0:
            fcsts = torch.vstack([fcst[:, -(1+(self.test_size-self.h)//self.step_size ):,:] for fcst in fcsts])
            fcsts = fcsts.numpy().flatten()
            fcsts = fcsts.reshape(-1, self.loss.outputsize_multiplier)
        else:
            fcsts = torch.vstack([fcst[:, -1:, :] for fcst in fcsts]).numpy().flatten()
            fcsts = fcsts.reshape(-1, self.loss.outputsize_multiplier)

        return fcsts
