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

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

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

from ..losses.pytorch import MAE
from ..common._base_windows import BaseWindows

# %% ../../nbs/models.mlp.ipynb 6
class MLP(BaseWindows):
    
    def __init__(self, 
                 input_size,
                 h,
                 step_size=1,
                 hidden_size=1024, 
                 num_layers=2, 
                 learning_rate=1e-3,
                 normalize=False,
                 loss=MAE(),
                 batch_size=32, 
                 num_workers_loader=0,
                 drop_last_loader=False,
                 random_seed=1,
                 **trainer_kwargs):
        
        # Inherit BaseWindows class
        super(MLP, self).__init__(h=h, 
                                  loss=loss,
                                  batch_size=batch_size,
                                  normalize=normalize,
                                  num_workers_loader=num_workers_loader,
                                  drop_last_loader=drop_last_loader,
                                  random_seed=random_seed,
                                  **trainer_kwargs)
        
        self.input_size = input_size
        self.step_size = step_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.learning_rate = learning_rate
        self.loss = loss
        
        # MultiLayer Perceptron
        layers = [nn.Linear(in_features=input_size, out_features=hidden_size)]
        for i in range(num_layers - 1):
            layers += [nn.Linear(in_features=hidden_size, out_features=hidden_size)]
        self.mlp = nn.ModuleList(layers)
        
        # Adapter with Loss dependent dimensions
        self.out = nn.Linear(in_features=hidden_size, 
                             out_features=h * self.loss.outputsize_multiplier)
        
    def forward(self, x, mask):
        y_pred = x
        for layer in self.mlp:
             y_pred = torch.relu(layer(y_pred))
        y_pred = self.out(y_pred)
        y_pred = self.loss.adapt_output(y_pred)
        return y_pred
