import torch.nn as nn

import torch.nn.functional as F

import torch.utils.data

from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig, BertForSequenceClassification

import math

import numpy as np

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

"""MLP baseline using sentence level eeg"""

# using sent level EEG, MLP baseline for sentiment

class BaselineMLPSentence(nn.Module):

    def __init__(self, input_dim = 840, hidden_dim = 128, output_dim = 3):

        super(BaselineMLPSentence, self).__init__()

        self.fc1 = nn.Linear(input_dim, hidden_dim) 

        self.relu1 = nn.ReLU()

        self.fc2 = nn.Linear(hidden_dim, hidden_dim)

        self.relu2 = nn.ReLU()

        self.fc3 = nn.Linear(hidden_dim, output_dim) # positive, negative, neutral  

        self.dropout = nn.Dropout(0.25)

    def forward(self, x):

        out = self.fc1(x)

        out = self.relu1(out)

        out = self.fc2(out)

        out = self.relu2(out)

        out = self.dropout(out)

        out = self.fc3(out)

        return out

"""bidirectional LSTM baseline using word level eeg"""

class BaselineLSTM(nn.Module):

    def __init__(self, input_dim = 840, hidden_dim = 256, output_dim = 3, num_layers = 1):

        super(BaselineLSTM, self).__init__()

        self.hidden_dim = hidden_dim

        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers = 1, batch_first = True, bidirectional = True)

        self.hidden2sentiment = nn.Linear(hidden_dim*2, output_dim)

    def forward(self, x_packed):

        # input: (N,seq_len,input_dim)

        # print(x_packed.data.size())

        lstm_out, _ = self.lstm(x_packed)

        last_hidden_state = pad_packed_sequence(lstm_out, batch_first = True)[0][:,-1,:]

        # print(last_hidden_state.size())

        out = self.hidden2sentiment(last_hidden_state)

        return out

""" Bert Baseline: Finetuning from a pretrained language model Bert"""

class NaiveFineTunePretrainedBert(nn.Module):

    def __init__(self, input_dim = 840, hidden_dim = 768, output_dim = 3, pretrained_checkpoint = None):

        super(NaiveFineTunePretrainedBert, self).__init__()

        # mapping hidden states dimensioin

        self.fc1 = nn.Linear(input_dim, hidden_dim)

        self.pretrained_Bert = BertForSequenceClassification.from_pretrained('bert-base-cased',num_labels=3)

        if pretrained_checkpoint is not None:

            self.pretrained_Bert.load_state_dict(torch.load(pretrained_checkpoint))

    def forward(self, input_embeddings_batch, input_masks_batch, labels):

        embedding = F.relu(self.fc1(input_embeddings_batch))

        out = self.pretrained_Bert(inputs_embeds = embedding, attention_mask = input_masks_batch, labels = labels, return_dict = True)

        return out

""" Finetuning from a pretrained language model BART, two step training"""

class FineTunePretrainedTwoStep(nn.Module):

    def __init__(self, pretrained_layers, in_feature = 840, d_model = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048):

        super(FineTunePretrainedTwoStep, self).__init__()

        self.pretrained_layers = pretrained_layers

        # additional transformer encoder, following BART paper about 

        self.additional_encoder_layer = nn.TransformerEncoderLayer(d_model=in_feature, nhead=additional_encoder_nhead,  dim_feedforward = additional_encoder_dim_feedforward, batch_first=True)

        self.additional_encoder = nn.TransformerEncoder(self.additional_encoder_layer, num_layers=6)

        # NOTE: add positional embedding?

        # print('[INFO]adding positional embedding')

        # self.positional_embedding = PositionalEncoding(in_feature)

        self.fc1 = nn.Linear(in_feature, d_model)

    def forward(self, input_embeddings_batch, input_masks_batch, input_masks_invert, labels):

        """input_embeddings_batch: batch_size*Seq_len*840"""

        """input_mask: 1 is not masked, 0 is masked"""

        """input_masks_invert: 1 is masked, 0 is not masked"""

        """labels: sentitment labels 0,1,2"""

        # NOTE: add positional embedding?

        # input_embeddings_batch = self.positional_embedding(input_embeddings_batch) 

        # use src_key_padding_masks

        encoded_embedding = self.additional_encoder(input_embeddings_batch, src_key_padding_mask = input_masks_invert) 

        # encoded_embedding = self.additional_encoder(input_embeddings_batch) 

        encoded_embedding = F.relu(self.fc1(encoded_embedding))

        out = self.pretrained_layers(inputs_embeds = encoded_embedding, attention_mask = input_masks_batch, return_dict = True, labels = labels)                    

        return out

""" Zero-shot sentiment discovery using a finetuned generation model and a sentiment model pretrained on text """

class ZeroShotSentimentDiscovery(nn.Module):

    def __init__(self, brain2text_translator, sentiment_classifier, translation_tokenizer, sentiment_tokenizer, device = 'cpu'):

        # only for inference

        super(ZeroShotSentimentDiscovery, self).__init__()

        self.brain2text_translator = brain2text_translator

        self.sentiment_classifier = sentiment_classifier

        self.translation_tokenizer = translation_tokenizer

        self.sentiment_tokenizer = sentiment_tokenizer

        self.device = device

    def forward(self, input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted, sentiment_labels):

        """input_embeddings_batch: batch_size*Seq_len*840"""

        """input_mask: 1 is not masked, 0 is masked"""

        """input_masks_invert: 1 is masked, 0 is not masked"""

        """labels: sentitment labels 0,1,2"""

        def logits2PredString(logits):

            probs = logits[0].softmax(dim = 1)

            # print('probs size:', probs.size())

            values, predictions = probs.topk(1)

            # print('predictions before squeeze:',predictions.size())

            predictions = torch.squeeze(predictions)

            predict_string = self.translation_tokenizer.decode(predictions)

            return predict_string

        # only works on batch is one

        assert input_embeddings_batch.size()[0] == 1

        seq2seqLMoutput = self.brain2text_translator(input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted)

        predict_string = logits2PredString(seq2seqLMoutput.logits)

        predict_string = predict_string.split('</s></s>')[0]

        predict_string = predict_string.replace('<s>','')

        print('predict string:', predict_string)

        re_tokenized = self.sentiment_tokenizer(predict_string, return_tensors='pt', return_attention_mask = True)

        input_ids = re_tokenized['input_ids'].to(self.device) # batch = 1

        attn_mask = re_tokenized['attention_mask'].to(self.device) # batch = 1

        out = self.sentiment_classifier(input_ids = input_ids, attention_mask = attn_mask, return_dict = True, labels = sentiment_labels)

        return out

""" Miscellaneous: jointly learn generation and classification (not working well) """

class BartClassificationHead(nn.Module):

    # from transformers: https://huggingface.co/transformers/_modules/transformers/models/bart/modeling_bart.html

    """Head for sentence-level classification tasks."""

    def __init__(

        self,

        input_dim: int,

        inner_dim: int,

        num_classes: int,

        pooler_dropout: float,

    ):

        super().__init__()

        self.dense = nn.Linear(input_dim, inner_dim)

        self.dropout = nn.Dropout(p=pooler_dropout)

        self.out_proj = nn.Linear(inner_dim, num_classes)

    def forward(self, hidden_states: torch.Tensor):

        hidden_states = self.dropout(hidden_states)

        hidden_states = self.dense(hidden_states)

        hidden_states = torch.tanh(hidden_states)

        hidden_states = self.dropout(hidden_states)

        hidden_states = self.out_proj(hidden_states)

        return hidden_states

class JointBrainTranslatorSentimentClassifier(nn.Module):

    def __init__(self, pretrained_layers, in_feature = 840, d_model = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048, num_labels = 3):

        super(JointBrainTranslatorSentimentClassifier, self).__init__()

        self.pretrained_generator = pretrained_layers

        # additional transformer encoder, following BART paper about 

        self.additional_encoder_layer = nn.TransformerEncoderLayer(d_model=in_feature, nhead=additional_encoder_nhead,  dim_feedforward = additional_encoder_dim_feedforward, batch_first=True)

        self.additional_encoder = nn.TransformerEncoder(self.additional_encoder_layer, num_layers=6)

        self.fc1 = nn.Linear(in_feature, d_model)

        self.num_labels = num_labels

        self.pooler = Pooler(d_model)

        self.classifier = BartClassificationHead(input_dim = d_model, inner_dim = d_model, num_classes = num_labels, pooler_dropout = pretrained_layers.config.classifier_dropout)

    def forward(self, input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted, sentiment_labels):

        """input_embeddings_batch: batch_size*Seq_len*840"""

        """input_mask: 1 is not masked, 0 is masked"""

        """input_masks_invert: 1 is masked, 0 is not masked"""

        # NOTE: add positional embedding?

        # input_embeddings_batch = self.positional_embedding(input_embeddings_batch) 

        # use src_key_padding_masks

        encoded_embedding = self.additional_encoder(input_embeddings_batch, src_key_padding_mask = input_masks_invert) 

        # encoded_embedding = self.additional_encoder(input_embeddings_batch) 

        encoded_embedding = F.relu(self.fc1(encoded_embedding))

        LMoutput = self.pretrained_generator(inputs_embeds = encoded_embedding, attention_mask = input_masks_batch, return_dict = True, labels = target_ids_batch_converted, output_hidden_states = True)                    

        hidden_states = LMoutput.decoder_hidden_states # N, seq_len, hidden_dim

        # print('hidden states len:', len(hidden_states))

        last_hidden_states = hidden_states[-1]

        # print('last hidden states size:', last_hidden_states.size())

        sentence_representation = self.pooler(last_hidden_states)

        classification_logits = self.classifier(sentence_representation) 

        loss_fct = nn.CrossEntropyLoss()

        classification_loss = loss_fct(classification_logits.view(-1, self.num_labels), sentiment_labels.view(-1))

        classification_output = {'loss':classification_loss,'logits':classification_logits}

        # print('successful one forward!!!!')

        return LMoutput, classification_output

""" helper modules """

# modified from BertPooler

class Pooler(nn.Module):

    def __init__(self, hidden_size):

        super().__init__()

        self.dense = nn.Linear(hidden_size, hidden_size)

        self.activation = nn.Tanh()

    def forward(self, hidden_states):

        # We "pool" the model by simply taking the hidden state corresponding

        # to the first token.

        first_token_tensor = hidden_states[:, 0]

        pooled_output = self.dense(first_token_tensor)

        pooled_output = self.activation(pooled_output)

        return pooled_output

# from https://pytorch.org/tutorials/beginner/transformer_tutorial.html

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout=0.1, max_len=5000):

        super(PositionalEncoding, self).__init__()

        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)

        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)

        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0).transpose(0, 1)

        self.register_buffer('pe', pe)

    def forward(self, x):

        # print('[DEBUG] input size:', x.size())

        # print('[DEBUG] positional embedding size:', self.pe.size())

        x = x + self.pe[:x.size(0), :]

        # print('[DEBUG] output x with pe size:', x.size())

        return self.dropout(x)