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

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

# %% ../../nbs/models.lstm.ipynb 5
import torch.nn as nn

from ..losses.pytorch import MAE
from ..common._base_recurrent import BaseRecurrent

# %% ../../nbs/models.lstm.ipynb 6
class LSTM(BaseRecurrent):
    def __init__(self,
                 input_size: int,
                 h: int,
                 state_hsize: int = 200, 
                 step_size: int = 1,
                 n_layers: int = 2,
                 bias: bool = True,
                 dropout: float = 0.,
                 learning_rate: float = 1e-3,
                 normalize: bool = True,
                 loss=MAE(),
                 batch_size=32, 
                 num_workers_loader=0,
                 drop_last_loader=False,
                 random_seed=1,
                 **trainer_kwargs):
        super(LSTM, self).__init__(
            loss=loss,
            batch_size=batch_size,
            num_workers_loader=num_workers_loader,
            drop_last_loader=drop_last_loader,
            random_seed=random_seed,
            **trainer_kwargs
        )

        # Architecture
        self.input_size = input_size
        self.h = h
        self.state_hsize = state_hsize
        self.step_size = step_size
        self.n_layers = n_layers
        self.bias = bias
        self.dropout = dropout

        # Optimization
        self.learning_rate = learning_rate
        self.loss = loss
        self.normalize = normalize
        self.random_seed = random_seed
        self.padder = nn.ConstantPad1d(padding=(0, self.h), value=0)

        # Instantiate model
        self.model = nn.LSTM(input_size=self.input_size,
                             hidden_size=self.state_hsize,
                             num_layers=self.n_layers,
                             bias=self.bias,
                             dropout=self.dropout,
                             batch_first=True)
        self.adapterW  = nn.Linear(self.state_hsize, self.h)

    def forward(self, insample_y, insample_mask):

        # LSTM forward
        insample_y, _ = self.model(insample_y)
        insample_y = self.adapterW(insample_y)
        
        return insample_y
