import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import torch.autograd.profiler as profiler
import termcolor
import numpy as np
import os

from datasets import load_dataset
from nltk.tokenize import wordpunct_tokenize

import torch
import torch.nn as nn
import torch.optim as optim

from transformers import BertConfig
from transformers.models.bert.modeling_bert import BertSelfAttention, BertSelfOutput, BertIntermediate, BertOutput

import logging

class AttentionPooling(nn.Module):
    def __init__(self, config):
        self.config = config
        super(AttentionPooling, self).__init__()
        self.att_fc1 = nn.Linear(config.hidden_size, config.hidden_size)
        self.att_fc2 = nn.Linear(config.hidden_size, 1)
        self.apply(self.init_weights)
        
    def init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()
            
                
    def forward(self, x, attn_mask=None):
        
        bz = x.shape[0]
        e = self.att_fc1(x)
        e = nn.Tanh()(e)
        alpha = self.att_fc2(e)
        alpha = torch.exp(alpha)
        if attn_mask is not None:
            alpha = alpha * attn_mask.unsqueeze(2)
        alpha = alpha / (torch.sum(alpha, dim=1, keepdim=True) + 1e-8)
        x = torch.bmm(x.permute(0, 2, 1), alpha)
        x = torch.reshape(x, (bz, -1))  
            
        return x

class BertAttention(nn.Module):
    def __init__(self, config):
        super(BertAttention, self).__init__()
        self.self = BertSelfAttention(config)
#         print("{} {}".format(self.self.all_head_size, config.hidden_size))
        self.output = BertSelfOutput(config)

    def forward(self, input_tensor, attention_mask):
        with profiler.record_function("BertSelfAttention-3"):
#             print("input_tensor {}, attention_mask {}".format(input_tensor.shape, attention_mask.shape))
            self_output = self.self(input_tensor, attention_mask.unsqueeze(1))[0]
#             print("self_output {}".format(self_output.shape))
            
        with profiler.record_function("BertSelfOutput-3"):
            attention_output = self.output(self_output, input_tensor)
        return attention_output

class BertformerLayer(nn.Module):
    def __init__(self, config):
        super(BertformerLayer, self).__init__()
        self.attention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden_states, attention_mask):
        with profiler.record_function("BertAttention-2"):
            attention_output = self.attention(hidden_states, attention_mask)
        with profiler.record_function("BertIntermediate-2"):
            intermediate_output = self.intermediate(attention_output)
        with profiler.record_function("BertOutput-2"):
            layer_output = self.output(intermediate_output, attention_output)
        return layer_output
    
class BertformerEncoder(nn.Module):
    def __init__(self, config, pooler_count=1):
        super(BertformerEncoder, self).__init__()
        self.config = config
        self.encoders = nn.ModuleList([BertformerLayer(config) for _ in range(config.num_hidden_layers)])
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # support multiple different poolers with shared bert encoder.
        self.poolers = nn.ModuleList()
        if config.pooler_type == 'weightpooler':
            for _ in range(pooler_count):
                self.poolers.append(AttentionPooling(config))
        logging.info(f"This model has {len(self.poolers)} poolers.")

        self.apply(self.init_weights)

    def init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, (nn.Embedding)) and module.padding_idx is not None:
                with torch.no_grad():
                    module.weight[module.padding_idx].fill_(0)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()

    def forward(self, 
                input_embs, 
                attention_mask, 
                pooler_index=0):
        #input_embs: batch_size, seq_len, emb_dim
        #attention_mask: batch_size, seq_len, emb_dim

        with profiler.record_function("PositionalEmbedding-1"):
            extended_attention_mask = attention_mask.unsqueeze(2)
            extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
            extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

            batch_size, seq_length, emb_dim = input_embs.shape
            position_ids = torch.arange(seq_length, dtype=torch.long, device=input_embs.device)
            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
            position_embeddings = self.position_embeddings(position_ids)

            embeddings = input_embs + position_embeddings

            embeddings = self.LayerNorm(embeddings)
            embeddings = self.dropout(embeddings)
#             print(embeddings.size())
            all_hidden_states = [embeddings]

        with profiler.record_function("BertformerLayer-1"):
            for i, layer_module in enumerate(self.encoders):
                layer_outputs = layer_module(all_hidden_states[-1], extended_attention_mask)
                all_hidden_states.append(layer_outputs)
            assert len(self.poolers) > pooler_index
            
        with profiler.record_function("Poolers-1"):
            output = self.poolers[pooler_index](all_hidden_states[-1], attention_mask)

        return output 


    
class BertModel(torch.nn.Module):

    def __init__(self,config, word_dict):
        super(BertModel, self).__init__()
        self.config = config
        self.dense_linear = nn.Linear(config.hidden_size,4)
        self.word_embedding = nn.Embedding(len(word_dict),256,padding_idx=0)
        self.bertformer_model = BertformerEncoder(config)
        self.criterion = nn.CrossEntropyLoss() 
        self.apply(self.init_weights)
        
    def init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, (nn.Embedding)) and module.padding_idx is not None:
                with torch.no_grad():
                    module.weight[module.padding_idx].fill_(0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()
    
    def forward(self,input_ids,targets):
        
        with profiler.record_function("Bert"):
            with profiler.record_function("Embedding-0"):
                mask=input_ids.bool().float()
                embds=self.word_embedding(input_ids)
    #             print("{} {} {}".format(mask.shape, input_ids.shape, embds.shape))

            with profiler.record_function("BertformerEncoder-0"):
                text_vec = self.bertformer_model(embds,mask)

            with profiler.record_function("Classifier-0"):
                score = self.dense_linear(text_vec)
                loss = self.criterion(score, targets) 

            return loss, score