import torch.nn as nn

import torch.nn.functional as F

import torch.utils.data

from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig, Text2TextGenerationPipeline

import math

import numpy as np

""" main architecture for open vocabulary EEG-To-Text decoding"""

class BrainTranslator(nn.Module):

    def __init__(self, pretrained_layers, in_feature=840, decoder_embedding_size=1024, additional_encoder_nhead=8,

                 additional_encoder_dim_feedforward=2048):

        super(BrainTranslator, self).__init__()

        self.pretrained = 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)

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

        # self.positional_embedding = PositionalEncoding(in_feature)

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

    def addin_forward(self, input_embeddings_batch, input_masks_invert):

        """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"""

        # 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))

        return encoded_embedding

    @torch.no_grad()

    def generate(

            self,

            input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted,

            generation_config=None,

            logits_processor=None,

            stopping_criteria=None,

            prefix_allowed_tokens_fn=None,

            synced_gpus=None,

            assistant_model=None,

            streamer=None,

            negative_prompt_ids=None,

            negative_prompt_attention_mask=None,

            **kwargs,

    ):

        encoded_embedding = self.addin_forward(input_embeddings_batch, input_masks_invert)

        output = self.pretrained.generate(

            inputs_embeds=encoded_embedding,

            attention_mask=input_masks_batch[:, :encoded_embedding.shape[1]],

            labels=target_ids_batch_converted,

            return_dict=True,

            generation_config=generation_config,

            logits_processor=logits_processor,

            stopping_criteria=stopping_criteria,

            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,

            synced_gpus=synced_gpus,

            assistant_model=assistant_model,

            streamer=streamer,

            negative_prompt_ids=negative_prompt_ids,

            negative_prompt_attention_mask=negative_prompt_attention_mask,

            **kwargs, )

        return output

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

        encoded_embedding = self.addin_forward(input_embeddings_batch, input_masks_invert)

        # print(f'forward:{input_embeddings_batch.shape,input_masks_batch.shape,input_masks_invert.shape,target_ids_batch_converted.shape,encoded_embedding.shape}')

        out = self.pretrained(inputs_embeds=encoded_embedding, attention_mask=input_masks_batch,

                              return_dict=True, labels=target_ids_batch_converted)

        return out

""" crippled open vocabulary EEG-To-Text decoding model w/o additional MTE encoder"""

class BrainTranslatorNaive(nn.Module):

    def __init__(self, pretrained_layers, in_feature=840, decoder_embedding_size=1024, additional_encoder_nhead=8,

                 additional_encoder_dim_feedforward=2048):

        super(BrainTranslatorNaive, self).__init__()

        '''no additional transformer encoder version'''

        self.pretrained = pretrained_layers

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

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

        """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"""

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

        out = self.pretrained(inputs_embeds=encoded_embedding, attention_mask=input_masks_batch, return_dict=True,

                              labels=target_ids_batch_converted)

        return out

""" 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)

""" Miscellaneous (not working well) """

class BrainTranslatorBert(nn.Module):

    def __init__(self, pretrained_layers, in_feature=840, hidden_size=768):

        super(BrainTranslatorBert, self).__init__()

        self.pretrained_Bert = pretrained_layers

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

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

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

        out = self.pretrained_Bert(inputs_embeds=embedding, attention_mask=input_masks_batch, labels=target_ids_batch,

                                   return_dict=True)

        return out

class EEG2BertMapping(nn.Module):

    def __init__(self, in_feature=840, hidden_size=512, out_feature=768):

        super(EEG2BertMapping, self).__init__()

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

        self.fc2 = nn.Linear(hidden_size, out_feature)

    def forward(self, x):

        out = F.relu(self.fc1(x))

        out = self.fc2(out)

        return out

class ContrastiveBrainTextEncoder(nn.Module):

    def __init__(self, pretrained_text_encoder, in_feature=840, eeg_encoder_nhead=8, eeg_encoder_dim_feedforward=2048,

                 embed_dim=768):

        super(ContrastiveBrainTextEncoder, self).__init__()

        # EEG Encoder

        self.positional_embedding = PositionalEncoding(in_feature)

        self.encoder_layer = nn.TransformerEncoderLayer(d_model=in_feature, nhead=eeg_encoder_nhead,

                                                        dim_feedforward=eeg_encoder_dim_feedforward, batch_first=True)

        self.EEG_Encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=6)

        self.EEG_pooler = Pooler(in_feature)

        self.ln_final = nn.LayerNorm(in_feature)  # to be considered

        # project to text embedding

        self.EEG_projection = nn.Parameter(torch.empty(in_feature, embed_dim))

        # Text Encoder

        self.TextEncoder = pretrained_text_encoder

        # learned temperature parameter

        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

    def forward(self, input_EEG_features, input_EEG_attn_mask, input_ids, input_text_attention_masks):

        # add positional embedding

        input_EEG_features = self.positional_embedding(input_EEG_features)

        # get EEG feature embedding

        EEG_hiddenstates = self.EEG_Encoder(input_EEG_features, src_key_padding_mask=input_EEG_attn_mask)

        EEG_hiddenstates = self.ln_final(EEG_hiddenstates)

        EEG_features = self.EEG_pooler(EEG_hiddenstates)  # [N, 840]

        # project to text embed size

        EEG_features = EEG_features @ self.EEG_projection  # [N, 768]

        # get text feature embedding

        Text_features = self.TextEncoder(input_ids=input_ids, attention_mask=input_text_attention_masks,

                                         return_dict=True).pooler_output  # [N, 768]

        # normalized features

        EEG_features = EEG_features / EEG_features.norm(dim=-1, keepdim=True)  # [N, 768]

        Text_features = Text_features / Text_features.norm(dim=-1, keepdim=True)  # [N, 768]

        # cosine similarity as logits

        logit_scale = self.logit_scale.exp()

        logits_per_EEG = logit_scale * EEG_features @ Text_features.t()  # [N, N]

        logits_per_text = logit_scale * Text_features @ EEG_features.t()  # [N, N]

        return logits_per_EEG, logits_per_text