# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/078_models.TransformerRNNPlus.ipynb.

# %% auto 0
__all__ = ['TransformerRNNPlus', 'TransformerLSTMPlus', 'TransformerGRUPlus']

# %% ../../nbs/078_models.TransformerRNNPlus.ipynb 3
import torch
from torch import nn
from torch.nn import TransformerEncoder, TransformerEncoderLayer
from collections import OrderedDict
from .layers import lin_nd_head

# %% ../../nbs/078_models.TransformerRNNPlus.ipynb 6
class _TransformerRNNEncoder(nn.Module):
    def __init__(self, 
        cell:nn.Module, # A RNN cell instance.
        c_in:int, # Number of channels in the input tensor.
        seq_len:int, # Number of time steps in the input tensor.
        d_model:int, # The number of expected features in the input.
        nhead:int, # Number of parallel attention heads (d_model will be split across nhead - each head will have dimension d_model // nhead).
        proj_dropout:float=0.1, # Dropout probability after the projection linear layer. Default: 0.1.
        num_encoder_layers:int=1, # Number of transformer layers in the encoder. Default: 1.
        dim_feedforward:int=2048, # The dimension of the feedforward network model. Default: 2048.
        dropout:float=0.1, # Transformer encoder layers dropout. Default: 0.1.
        num_rnn_layers:int=1, # Number of RNN layers in the encoder. Default: 1.
        bidirectional:bool=True, # If True, becomes a bidirectional RNN. Default: True.
        ):
        super().__init__()

        # projection layer
        self.proj_linear = nn.Linear(c_in, d_model)
        self.proj_dropout = nn.Dropout(proj_dropout)

        # transformer encoder layers
        self.num_encoder_layers = num_encoder_layers
        dim_feedforward = dim_feedforward or d_model
        self.enc_layers = nn.ModuleList([TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True, ) for _ in range(num_encoder_layers)])

        # rnn layers
        self.num_rnn_layers = num_rnn_layers
        self.rnn_layers = nn.ModuleList([cell(d_model * (1 + bidirectional) ** i, d_model * (1 + bidirectional) ** i, bidirectional=bidirectional) for i in range(num_rnn_layers)])
        self.seq_len = seq_len
        self.pos_encoding = nn.Parameter(torch.randn(1, self.seq_len, d_model) * 0.02)

    def forward(self, x): # (batch_size, c_in, seq_len), Example shape (4, 54, 864) 
        
        x = x.swapaxes(1, 2) # (batch_size, seq_len, c_in), Example shape (4, 864, 54)
        batch_size = x.shape[0]

        # projection layer
        x = self.proj_linear(x) # (batch_size, seq_len, d_model), Example shape (4, 864, 320)
        x = x + self.pos_encoding.repeat(batch_size, 1, 1)
        x = self.proj_dropout(x)

        # transformer encoder layers
        for i in range(self.num_encoder_layers): 
            x = self.enc_layers[i](x.transpose(0, 1)).transpose(0, 1) # (batch_size, seq_len, d_model), Example shape (4, 864, 320)
        
        # rnn layers
        for i in range(self.num_rnn_layers): 
            x, _ = self.rnn_layers[i](x) # (batch_size, seq_len, CFG['fog_model_dim']*2), Example shape (4, 864, 640)
        
        x = x.swapaxes(1, 2)
        return x

# %% ../../nbs/078_models.TransformerRNNPlus.ipynb 8
class _TransformerRNNPlus(nn.Sequential):
    def __init__(self, 
        c_in:int, # Number of channels in the input tensor.
        c_out:int, # Number of output channels.
        seq_len:int, # Number of time steps in the input tensor.
        d:tuple=None, # int or tuple with shape of the output tensor
        d_model:int=128, # Total dimension of the model.
        nhead:int=16, # Number of parallel attention heads (d_model will be split across nhead - each head will have dimension d_model // nhead).
        proj_dropout:float=0.1, # Dropout probability after the first linear layer. Default: 0.1.
        num_encoder_layers:int=1, # Number of transformer encoder layers. Default: 1.
        dim_feedforward:int=2048, # The dimension of the feedforward network model. Default: 2048.
        dropout:float=0.1, # Transformer encoder layers dropout. Default: 0.1.
        num_rnn_layers:int=1, # Number of RNN layers in the encoder. Default: 1.
        bidirectional:bool=True, # If True, becomes a bidirectional RNN. Default: True.
        custom_head=None, # Custom head that will be applied to the model. If None, a head with `c_out` outputs will be used. Default: None.
        **kwargs
        ):

        backbone = _TransformerRNNEncoder(cell=self._cell, c_in=c_in, seq_len=seq_len, proj_dropout=proj_dropout, 
                                          d_model=d_model, nhead=nhead, num_encoder_layers=num_encoder_layers, dim_feedforward=dim_feedforward, dropout=dropout, 
                                          num_rnn_layers=num_rnn_layers, bidirectional=bidirectional)
        self.head_nf = d_model * ((1 + bidirectional) ** (num_rnn_layers))
        if custom_head:
            if isinstance(custom_head, nn.Module): head = custom_head
            else: head = custom_head(self.head_nf, c_out, seq_len, d=d, **kwargs)
        else:
            head = lin_nd_head(self.head_nf, c_out, seq_len=seq_len, d=d)
        super().__init__(OrderedDict([('backbone', backbone), ('head', head)]))

class TransformerRNNPlus(_TransformerRNNPlus):
    _cell = nn.RNN

class TransformerLSTMPlus(_TransformerRNNPlus):
    _cell = nn.LSTM

class TransformerGRUPlus(_TransformerRNNPlus):
    _cell = nn.GRU
