# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/075b_models.RNNAttentionPlus.ipynb.

# %% auto 0
__all__ = ['RNNAttentionPlus', 'LSTMAttentionPlus', 'GRUAttentionPlus']

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 6
from ..imports import *
from ..utils import *
from .layers import *
from .utils import *

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 8
class _ScaledDotProductAttention(Module):
    def __init__(self, d_k:int): self.d_k = d_k
    def forward(self, q:Tensor, k:Tensor, v:Tensor, mask:Optional[Tensor]=None):

        # MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
        scores = torch.matmul(q, k)                                         # scores : [bs x n_heads x q_len x q_len]
        
        # Scale
        scores = scores / (self.d_k ** 0.5)
        
        # Mask (optional)
        if mask is not None: scores.masked_fill_(mask, -1e9)
        
        # SoftMax
        attn = F.softmax(scores, dim=-1)                                    # attn   : [bs x n_heads x q_len x q_len]
        
        # MatMul (attn, v)
        context = torch.matmul(attn, v)                                     # context: [bs x n_heads x q_len x d_v]
        
        return context, attn

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 9
class _MultiHeadAttention(Module):
    def __init__(self, d_model:int, n_heads:int, d_k:int, d_v:int):
        r"""
        Input shape:  Q, K, V:[batch_size (bs) x q_len x d_model], mask:[q_len x q_len]
        """
        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
        
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        
        self.W_O = nn.Linear(n_heads * d_v, d_model, bias=False)

    def forward(self, Q:Tensor, K:Tensor, V:Tensor, mask:Optional[Tensor]=None):
        
        bs = Q.size(0)

        # Linear (+ split in multiple heads)
        q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,2)       # q_s    : [bs x n_heads x q_len x d_k]
        k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0,2,3,1)     # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
        v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1,2)       # v_s    : [bs x n_heads x q_len x d_v]

        # Scaled Dot-Product Attention (multiple heads)
        context, attn = _ScaledDotProductAttention(self.d_k)(q_s, k_s, v_s)          # context: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len]

        # Concat
        context = context.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * self.d_v) # context: [bs x q_len x n_heads * d_v]

        # Linear
        output = self.W_O(context)                                                  # context: [bs x q_len x d_model]
        
        return output, attn

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 11
def get_activation_fn(activation):
    if activation == "relu": return nn.ReLU()
    elif activation == "gelu": return nn.GELU()
    else: return activation()
#         raise ValueError(f'{activation} is not available. You can use "relu" or "gelu"')

class _TSTEncoderLayer(Module):
    def __init__(self, q_len:int, d_model:int, n_heads:int, d_k:Optional[int]=None, d_v:Optional[int]=None, d_ff:int=256, dropout:float=0.1, 
                 activation:str="gelu"):

        assert d_model // n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
        d_k = ifnone(d_k, d_model // n_heads)
        d_v = ifnone(d_v, d_model // n_heads)

        # Multi-Head attention
        self.self_attn = _MultiHeadAttention(d_model, n_heads, d_k, d_v)

        # Add & Norm
        self.dropout_attn = nn.Dropout(dropout)
        self.batchnorm_attn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))

        # Position-wise Feed-Forward
        self.ff = nn.Sequential(nn.Linear(d_model, d_ff), 
                                get_activation_fn(activation), 
                                nn.Dropout(dropout), 
                                nn.Linear(d_ff, d_model))

        # Add & Norm
        self.dropout_ffn = nn.Dropout(dropout)
        self.batchnorm_ffn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))

    def forward(self, src:Tensor, mask:Optional[Tensor]=None) -> Tensor:

        # Multi-Head attention sublayer
        ## Multi-Head attention
        src2, attn = self.self_attn(src, src, src, mask=mask)
        ## Add & Norm
        src = src + self.dropout_attn(src2) # Add: residual connection with residual dropout
        src = self.batchnorm_attn(src)      # Norm: batchnorm 

        # Feed-forward sublayer
        ## Position-wise Feed-Forward
        src2 = self.ff(src)
        ## Add & Norm
        src = src + self.dropout_ffn(src2) # Add: residual connection with residual dropout
        src = self.batchnorm_ffn(src) # Norm: batchnorm

        return src

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 13
class _TSTEncoder(Module):
    def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=None, dropout=0.1, activation='gelu', n_layers=1):
        
        self.layers = nn.ModuleList([_TSTEncoderLayer(q_len, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, dropout=dropout, 
                                                            activation=activation) for i in range(n_layers)])

    def forward(self, src):
        output = src
        for mod in self.layers: output = mod(output)
        return output

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 14
class _RNNAttentionPlus_Backbone(nn.Module):
    def __init__(self, cell, c_in:int, seq_len:int, d:tuple=None, hidden_size=128, rnn_layers=1, bias=True, rnn_dropout=0, bidirectional=False, 
                 encoder_layers:int=3, n_heads:int=16, d_k:Optional[int]=None, d_v:Optional[int]=None,  
                 d_ff:int=256, encoder_dropout:float=0.1, act:str="gelu"):
        
        super().__init__()
        
        # RNN
        self.rnn = cell(c_in, hidden_size, num_layers=rnn_layers, bias=bias, batch_first=True, dropout=rnn_dropout, bidirectional=bidirectional)
        
        # Attention Encoder
        d_model = hidden_size * (1 + bidirectional)
        self.encoder = _TSTEncoder(seq_len, d_model, n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, dropout=encoder_dropout, activation=act, n_layers=encoder_layers)

    def forward(self, x:Tensor, mask:Optional[Tensor]=None) -> Tensor:  # x: [bs x nvars x q_len]

        # RNN
        x = x.transpose(2,1)    # [bs x nvars x q_len] --> [bs x q_len x nvars]
        output, _ = self.rnn(x) # output from all sequence steps: [bs x q_len x hidden_size * (1 + bidirectional)]

        # Attention Encoder
        z = self.encoder(output)                                             # z: [bs x q_len x d_model]
        z = z.transpose(2,1).contiguous()                               # z: [bs x d_model x q_len]    

        return z

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 15
class _RNNAttentionPlus_Base(nn.Sequential):
    def __init__(self, 
        c_in:int, # the number of features (aka variables, dimensions, channels) in the time series dataset.
        c_out:int, # the number of target classes.
        seq_len:int, # number of time steps in the time series. 
        d:tuple=None,  # output shape (excluding batch dimension).
        hidden_size:int=128, # the number of features in the hidden state h
        rnn_layers:int=1, # the number of recurrent layers of the RNN model. 
        bias:bool=True, # If ``False``, then the layer does not use bias weights `b_ih` and `b_hh`.
        rnn_dropout:float=0, # rnn dropout applied to the output of each RNN layer except the last layer.
        bidirectional=False, # If ``True``, becomes a bidirectional RNN. Default: ``False``
        encoder_layers:int=3, # the number of sub-encoder-layers in the encoder.
        n_heads:int=16, # parallel attention heads.
        d_k:Optional[int]=None, # size of the learned linear projection of queries and keys in the MHA.
        d_v:Optional[int]=None, # size of the learned linear projection of values in the MHA.
        d_ff:int=256, # the dimension of the feedforward network model.
        encoder_dropout:float=0.1, # amount of residual dropout applied in the encoder.
        act:str="gelu", # the activation function of intermediate layer, relu or gelu.
        fc_dropout:float=0., # dropout applied to the final fully connected layer.
        y_range:Optional[tuple]=None, # range of possible y values (used in regression tasks).
        custom_head=None, # custom head that will be applied to the model head (optional).
        use_bn:bool=True, # indicates if batchnorm will be applied to the model head.
        flatten:bool=True, # this will flatten the output of the encoder before applying the head if True.
        ):
        "RNNAttention is a model based on RNN and TST (Time Series Transformer). TST is a Transformer that takes continuous time series as inputs."

        # Backbone
        backbone = _RNNAttentionPlus_Backbone(self._cell, c_in, seq_len, hidden_size=hidden_size, rnn_layers=rnn_layers, bias=bias, 
                                              rnn_dropout=rnn_dropout, bidirectional=bidirectional, 
                                              encoder_layers=encoder_layers, n_heads=n_heads, d_k=d_k, d_v=d_v,  
                                              d_ff=d_ff, encoder_dropout=encoder_dropout, act=act)

        # Head
        d_model = hidden_size * (1 + bidirectional)
        self.head_nf = seq_len * d_model
        if custom_head is not None: 
            if isinstance(custom_head, nn.Module): head = custom_head
            else: head = custom_head(d_model, c_out, seq_len)
        elif d is not None:
            head = lin_nd_head(d_model, c_out, seq_len=seq_len, d=d, use_bn=use_bn, fc_dropout=fc_dropout, flatten=flatten)
        else:
            head = self.create_head(self.head_nf, c_out, act=act, fc_dropout=fc_dropout, y_range=y_range)

        layers = OrderedDict([('backbone', backbone), ('head', head)])
        super().__init__(layers)
        

    def create_head(self, nf, c_out, act="gelu", fc_dropout=0., y_range=None, **kwargs):
        layers = [get_activation_fn(act), Flatten()]
        if fc_dropout: layers += [nn.Dropout(fc_dropout)]
        layers += [nn.Linear(nf, c_out)]
        if y_range: layers += [SigmoidRange(*y_range)]
        return nn.Sequential(*layers)    

# %% ../../nbs/075b_models.RNNAttentionPlus.ipynb 16
class RNNAttentionPlus(_RNNAttentionPlus_Base):
    _cell = nn.RNN
    
class LSTMAttentionPlus(_RNNAttentionPlus_Base):
    _cell = nn.LSTM
    
class GRUAttentionPlus(_RNNAttentionPlus_Base):
    _cell = nn.GRU
