# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/122_models.TabFusionTransformer.ipynb.

# %% auto 0
__all__ = ['ifnone', 'Sequential', 'TabFusionBackbone', 'TabFusionTransformer', 'TSTabFusionTransformer']

# %% ../../nbs/122_models.TabFusionTransformer.ipynb 4
# This is a modified Pytorch implementation based on TabTransformer created by Ignacio Oguiza (timeseriesAI@gmail.com):
# Huang, X., Khetan, A., Cvitkovic, M., & Karnin, Z. (2020). 
# TabTransformer: Tabular Data Modeling Using Contextual Embeddings. 
# arXiv preprint https://arxiv.org/pdf/2012.06678
# Official repo: https://github.com/awslabs/autogluon/tree/master/tabular/src/autogluon/tabular/models/tab_transformer

import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict


def ifnone(a, b):
    # From fastai.fastcore
    "`b` if `a` is None else `a`"
    return b if a is None else a


class _Flatten(nn.Module):
    def __init__(self, full=False):
        super().__init__()
        self.full = full
    def forward(self, x):
        return x.view(-1) if self.full else x.view(x.size(0), -1)


class Sequential(nn.Sequential):
    """Class that allows you to pass one or multiple inputs"""
    def forward(self, *x):
        for i, module in enumerate(self._modules.values()):
            x = module(*x) if isinstance(x, (list, tuple)) else module(x)
        return x


class _MLP(nn.Module):
    def __init__(self, dims, bn=False, act=None, skip=False, dropout=0., bn_final=False):
        super().__init__()
        dims_pairs = list(zip(dims[:-1], dims[1:]))
        layers = []
        for i, (dim_in, dim_out) in enumerate(dims_pairs):
            is_last = i >= (len(dims) - 2)
            if bn and (not is_last or bn_final): layers.append(nn.BatchNorm1d(dim_in))
            if dropout and not is_last:
                layers.append(nn.Dropout(dropout))
            layers.append(nn.Linear(dim_in, dim_out))
            if is_last: break
            layers.append(ifnone(act, nn.ReLU()))
        self.mlp = nn.Sequential(*layers)
        self.shortcut = nn.Linear(dims[0], dims[-1]) if skip else None

    def forward(self, x):
        if self.shortcut is not None: 
            return self.mlp(x) + self.shortcut(x)
        else:
            return self.mlp(x)
        

class _ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k:int, res_attention:bool=False): 
        super().__init__()
        self.d_k,self.res_attention = d_k,res_attention
        
    def forward(self, q, k, v, prev=None, attn_mask=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)

        # Attention mask (optional)
        if attn_mask is not None:                                     # mask with shape [q_len x q_len]
            if attn_mask.dtype == torch.bool:
                scores.masked_fill_(attn_mask, float('-inf'))
            else:
                scores += attn_mask

        # SoftMax
        if prev is not None: scores = scores + prev

        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]

        if self.res_attention: return context, attn, scores
        else: return context, attn


class _MultiheadAttention(nn.Module):
    def __init__(self, d_model:int, n_heads:int, d_k:int, d_v:int, res_attention:bool=False):
        """Input shape:  Q, K, V:[batch_size (bs) x q_len x d_model], mask:[q_len x q_len]"""
        super().__init__()
        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)

        self.res_attention = res_attention

        # Scaled Dot-Product Attention (multiple heads)
        if self.res_attention:
            self.sdp_attn = _ScaledDotProductAttention(self.d_k, self.res_attention)
        else:
            self.sdp_attn = _ScaledDotProductAttention(self.d_k)

        
    def forward(self, Q, K, V, prev=None, attn_mask=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)
        if self.res_attention:
            context, attn, scores = self.sdp_attn(q_s, k_s, v_s, prev=prev, attn_mask=attn_mask)
        else:
            context, attn = self.sdp_attn(q_s, k_s, v_s, attn_mask=attn_mask)
        # 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]

        if self.res_attention: return output, attn, scores
        else: return output, attn                                                            # output: [bs x q_len x d_model]


class _TabFusionEncoderLayer(nn.Module):
    def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=None, 
                 res_dropout=0.1, activation="gelu", res_attention=False):

        super().__init__()
        assert not 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)
        d_ff = ifnone(d_ff, d_model * 4)

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

        # Add & Norm
        self.dropout_attn = nn.Dropout(res_dropout)
        self.layernorm_attn = nn.LayerNorm(d_model)

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

        # Add & Norm
        self.dropout_ffn = nn.Dropout(res_dropout)
        self.layernorm_ffn = nn.LayerNorm(d_model)

    def forward(self, src, prev=None, attn_mask=None):

        # Multi-Head attention sublayer
        ## Multi-Head attention
        if self.res_attention:
            src2, attn, scores = self.self_attn(src, src, src, prev, attn_mask=attn_mask)
        else:
            src2, attn = self.self_attn(src, src, src, attn_mask=attn_mask)
        self.attn = attn
        ## Add & Norm
        src = src + self.dropout_attn(src2) # Add: residual connection with residual dropout
        src = self.layernorm_attn(src) # Norm: layernorm 

        # 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.layernorm_ffn(src) # Norm: layernorm

        if self.res_attention:
            return src, scores
        else:
            return src

    def _get_activation_fn(self, activation):
        if callable(activation): return activation()
        elif activation.lower() == "relu": return nn.ReLU()
        elif activation.lower() == "gelu": return nn.GELU()
        elif activation.lower() == "mish": return Mish()
        raise ValueError(f'{activation} is not available. You can use "relu", "gelu", "mish" or a callable')


class _TabFusionEncoder(nn.Module):
    def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=None, res_dropout=0.1, activation='gelu', res_attention=False, n_layers=1):
        super().__init__()
        self.layers = nn.ModuleList([_TabFusionEncoderLayer(q_len, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, res_dropout=res_dropout, 
                                                            activation=activation, res_attention=res_attention) for i in range(n_layers)])
        self.res_attention = res_attention

    def forward(self, src, attn_mask=None):
        output = src
        scores = None
        if self.res_attention:
            for mod in self.layers: output, scores = mod(output, prev=scores, attn_mask=attn_mask)
            return output
        else:
            for mod in self.layers: output = mod(output, attn_mask=attn_mask)
            return output


class TabFusionBackbone(nn.Module):
    def __init__(self, classes, cont_names, d_model=32, n_layers=6, n_heads=8, d_k=None, d_v=None, d_ff=None, init=True,
                 res_attention=True, attention_act='gelu', res_dropout=0.):

        super().__init__()
        
        # Categorical
        n_cat = len(classes)
        n_classes = [len(v) for v in classes.values()]
        self.n_emb = sum(n_classes)
        self.embeds = nn.ModuleList([nn.Embedding(ni, d_model) for ni in n_classes])
        
        # Continuous
        n_cont = len(cont_names)
        self.n_cont = n_cont
        self.conv = nn.Conv1d(1, d_model, 1)
        if init: nn.init.kaiming_normal_(self.conv.weight)

        # Transformer
        self.res_drop = nn.Dropout(res_dropout) if res_dropout else None
        self.pos_enc = nn.Parameter(torch.zeros(1, (n_cat  + n_cont), d_model))
        self.transformer = _TabFusionEncoder(n_cat + n_cont, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, res_dropout=res_dropout, 
                                             activation=attention_act, res_attention=res_attention, n_layers=n_layers)


    def forward(self, x_cat, x_cont=None):
        
        # Input encoding
        if self.n_emb != 0:
            x = [e(x_cat[:,i]).unsqueeze(1) for i,e in enumerate(self.embeds)]
            x = torch.cat(x, 1)
        if self.n_cont != 0:
            x_cont = self.conv(x_cont.unsqueeze(1)).transpose(1,2)
            x = torch.cat([x, x_cont], 1) if self.n_emb != 0 else x_cont

        # Transformer
        x += self.pos_enc
        if self.res_drop is not None: x = self.res_drop(x)
        x = self.transformer(x)

        return x


class TabFusionTransformer(Sequential):
    def __init__(self, classes, cont_names, c_out, 
                 d_model=32, n_layers=6, n_heads=8, d_k=None, d_v=None, d_ff=None, res_attention=True, attention_act='gelu', res_dropout=0.,
                 fc_mults=(4, 2), fc_dropout=0., fc_act=None, fc_skip=False, fc_bn=False, bn_final=False, init=True):

        super().__init__()
        
        # Backbone
        backbone = TabFusionBackbone(classes, cont_names, d_model=d_model, n_layers=n_layers, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, init=init,
                                     res_attention=res_attention, attention_act=attention_act, res_dropout=res_dropout)
        
        # Head
        mlp_input_size = (d_model * (len(classes)  + len(cont_names)))
        hidden_dimensions = list(map(lambda t: int(mlp_input_size * t), fc_mults))
        all_dimensions = [mlp_input_size, *hidden_dimensions, c_out]
        self.head_nf = mlp_input_size
        head = nn.Sequential(*[_Flatten(), _MLP(all_dimensions, act=fc_act, skip=fc_skip, bn=fc_bn, dropout=fc_dropout, bn_final=bn_final)])

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

# %% ../../nbs/122_models.TabFusionTransformer.ipynb 7
class TSTabFusionTransformer(nn.Module):
    def __init__(self, c_in, c_out, seq_len, classes, cont_names, 
                 d_model=32, n_layers=6, n_heads=8, d_k=None, d_v=None, d_ff=None, res_attention=True, attention_act='gelu', res_dropout=0., 
                 fc_mults=(1, .5), fc_dropout=0., fc_act=None, fc_skip=False, fc_bn=False, bn_final=False, init=True):

        super().__init__()
        
        # Time series
        self.W_P = nn.Conv1d(c_in, d_model, 1)
        
        # Categorical
        n_cat = len(classes)
        n_classes = [len(v) for v in classes.values()]
        self.n_emb = sum(n_classes)
        self.embeds = nn.ModuleList([nn.Embedding(ni, d_model) for ni in n_classes])
        
        # Continuous
        n_cont = len(cont_names)
        self.n_cont = n_cont
        self.conv = nn.Conv1d(1, d_model, 1)
        if init: nn.init.kaiming_normal_(self.conv.weight)

        # Transformer
        self.res_drop = nn.Dropout(res_dropout) if res_dropout else None
        self.pos_enc = nn.Parameter(torch.zeros(1, (n_cat  + n_cont + seq_len), d_model))
        self.transformer = _TabFusionEncoder(n_cat + n_cont, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, res_dropout=res_dropout, 
                                             activation=attention_act, res_attention=res_attention, n_layers=n_layers)
        
        # Head
        mlp_input_size = (d_model * (n_cat  + n_cont + seq_len))
        hidden_dimensions = list(map(lambda t: int(mlp_input_size * t), fc_mults))
        all_dimensions = [mlp_input_size, *hidden_dimensions, c_out]
        self.head_nf = mlp_input_size
        self.head = nn.Sequential(*[_Flatten(), _MLP(all_dimensions, act=fc_act, skip=fc_skip, bn=fc_bn, dropout=fc_dropout, bn_final=bn_final)])

    def forward(self, x):
        x_ts, (x_cat, x_cont) = x
        
        # Time series
        x = self.W_P(x_ts).transpose(1,2)
        
        # Input encoding
        if self.n_emb != 0:
            x_cat = [e(x_cat[:,i]).unsqueeze(1) for i,e in enumerate(self.embeds)]
            x_cat = torch.cat(x_cat, 1)
            x = torch.cat([x, x_cat], 1)
        if self.n_cont != 0:
            x_cont = self.conv(x_cont.unsqueeze(1)).transpose(1,2)
            x = torch.cat([x, x_cont], 1)

        # Transformer
        x += self.pos_enc
        if self.res_drop is not None: x = self.res_drop(x)
        x = self.transformer(x)

        # Head
        x = self.head(x)
        return x
