import numpy as np
import torch
import torch.nn as nn

from transformers.models.distilbert.modeling_distilbert import Transformer as T
from .pooling import create_pool1d_layer


class Config:

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)


class Transformer(nn.Module):
    """
    If predict_sequence is True, then the model will predict an output
    for each element in the sequence. If False, then the model will
    predict a single output for the sequence. 
    
    e.g., classifying each image in a CT scan vs the entire CT scan
    """
    def __init__(self,
                 num_classes,
                 embedding_dim=512,
                 hidden_dim=1024,
                 n_layers=4,
                 n_heads=16,
                 dropout=0.2,
                 attention_dropout=0.1,
                 output_attentions=False,
                 activation='gelu',
                 output_hidden_states=False,
                 chunk_size_feed_forward=0,
                 predict_sequence=True,
                 pool=None
    ):
        super().__init__()
        config = Config(**{
                'dim': embedding_dim,
                'hidden_dim': hidden_dim,
                'n_layers': n_layers,
                'n_heads': n_heads,
                'dropout': dropout,
                'attention_dropout': attention_dropout,
                'output_attentions': output_attentions,
                'activation': activation,
                'output_hidden_states': output_hidden_states,
                'chunk_size_feed_forward': chunk_size_feed_forward
            })

        self.transformer = T(config)
        self.predict_sequence = predict_sequence

        if not predict_sequence:
            if isinstance(pool, str):
                self.pool_layer = create_pool1d_layer(pool)
                if pool == "catavgmax":
                    embedding_dim *= 2
            else:
                self.pool_layer = nn.Identity()

        self.classifier = nn.Linear(embedding_dim, num_classes)

    def extract_features(self, x):
        x, mask = x
        x = self.transformer(x, attn_mask=mask, head_mask=[None]*x.size(1))
        x = x[0]

        if not self.predict_sequence:
            if isinstance(self.pool_layer, nn.Identity):
                # Just take the last vector in the sequence
                x = x[:, 0]
            else:
                x = self.pool_layer(x.transpose(-1, -2))

        return x

    def classify(self, x):

        if not self.predict_sequence:
            if isinstance(self.pool_layer, nn.Identity):
                # Just take the last vector in the sequence
                x = x[:, 0]
            else:
                x = self.pool_layer(x.transpose(-1, -2))

        out = self.classifier(x)

        if self.classifier.out_features == 1:
            return out[..., 0]
        else:
            return out

    def forward_tr(self, x, mask):
        output = self.transformer(x, attn_mask=mask, head_mask=[None]*x.size(1))
        return self.classify(output[0])

    def forward(self, x):
        x, mask = x
        return self.forward_tr(x, mask)


class DualTransformer(nn.Module):
    """
    Essentially the same as above except predicts both sequence and study labels.
    """
    def __init__(self,
                 num_classes,
                 embedding_dim=512,
                 hidden_dim=1024,
                 n_layers=4,
                 n_heads=16,
                 dropout=0.2,
                 attention_dropout=0.1,
                 output_attentions=False,
                 activation='gelu',
                 output_hidden_states=False,
                 chunk_size_feed_forward=0,
                 pool=None
    ):
        super().__init__()
        config = Config(**{
                'dim': embedding_dim,
                'hidden_dim': hidden_dim,
                'n_layers': n_layers,
                'n_heads': n_heads,
                'dropout': dropout,
                'attention_dropout': attention_dropout,
                'output_attentions': output_attentions,
                'activation': activation,
                'output_hidden_states': output_hidden_states,
                'chunk_size_feed_forward': chunk_size_feed_forward
            })

        self.transformer = T(config)

        self.pool_layer = nn.Identity()

        self.classifier1 = nn.Linear(embedding_dim, num_classes)
        self.classifier2 = nn.Linear(embedding_dim, num_classes)

    def classify(self, x):

        if isinstance(self.pool_layer, nn.Identity):
            # Just take the last vector in the sequence
            x_summ = x[:, 0]
        else:
            x_summ = self.pool_layer(x.transpose(-1, -2))

        # Element-wise labels
        out1 = self.classifier1(x)[:, :, 0]
        # Single label for whole sequence
        out2 = self.classifier2(x_summ)
        out = torch.cat([out1, out2], dim=1)
        return out

    def forward_tr(self, x, mask):
        output = self.transformer(x, attn_mask=mask, head_mask=[None]*x.size(1))
        return self.classify(output[0])

    def forward(self, x):
        x, mask = x
        return self.forward_tr(x, mask)


class DualTransformerV2(nn.Module):
    """
    More complicated variant of DualTransformer.
    Returns tuple of: (element-wise prediction, sequence prediction)
    """
    def __init__(self,
                 num_seq_classes,
                 num_classes,
                 embedding_dim=512,
                 hidden_dim=1024,
                 n_layers=4,
                 n_heads=16,
                 dropout=0.2,
                 attention_dropout=0.1,
                 output_attentions=False,
                 activation='gelu',
                 output_hidden_states=False,
                 chunk_size_feed_forward=0,
                 pool=None
        ):
        super().__init__()
        config = Config(**{
                'dim': embedding_dim,
                'hidden_dim': hidden_dim,
                'n_layers': n_layers,
                'n_heads': n_heads,
                'dropout': dropout,
                'attention_dropout': attention_dropout,
                'output_attentions': output_attentions,
                'activation': activation,
                'output_hidden_states': output_hidden_states,
                'chunk_size_feed_forward': chunk_size_feed_forward
            })

        self.transformer = T(config)

        self.pool_layer = nn.Identity()

        self.classifier1 = nn.Linear(embedding_dim, num_seq_classes)
        self.classifier2 = nn.Linear(embedding_dim, num_classes)

    def classify(self, x):

        if isinstance(self.pool_layer, nn.Identity):
            # Just take the last vector in the sequence
            x_summ = x[:, 0]
        else:
            x_summ = self.pool_layer(x.transpose(-1, -2))

        # Element-wise labels
        out1 = self.classifier1(x)
        if self.classifier1.out_features == 1:
            out1 = out1[:, :, 0]
        # Single label for whole sequence
        out2 = self.classifier2(x_summ)
        return out1, out2

    def forward_tr(self, x, mask):
        output = self.transformer(x, attn_mask=mask, head_mask=[None]*x.size(1))
        return self.classify(output[0])

    def forward(self, x):
        x, mask = x
        return self.forward_tr(x, mask)