'''
Seq2VecEncoders for encoding mentions and entities.
'''
import torch
import torch.nn as nn
from allennlp.modules.seq2vec_encoders import Seq2VecEncoder, PytorchSeq2VecWrapper, BagOfEmbeddingsEncoder
from allennlp.modules.seq2vec_encoders import BertPooler
from allennlp.nn.util import get_text_field_mask, add_positional_features
import numpy as np
from tqdm import tqdm
from allennlp.nn import util as nn_util
from allennlp.data.iterators import BasicIterator
import pdb
from allennlp.models import Model

# Custom imports
import os
from utils import pickle_load_object, pickle_save_object

class Pooler_for_title_and_desc(Seq2VecEncoder):
    def __init__(self, args, word_embedder):
        super(Pooler_for_title_and_desc, self).__init__()
        self.args = args
        self.huggingface_nameloader()
        self.bertpooler_sec2vec = BertPooler(pretrained_model=self.bert_weight_filepath)
        self.word_embedder = word_embedder
        self.word_embedding_dropout = nn.Dropout(self.args.word_embedding_dropout)

        self.linear_for_entity_encoding = nn.Linear(self.bertpooler_sec2vec.get_output_dim(),
                                                     self.bertpooler_sec2vec.get_output_dim())

        self.linear_for_dimentionReduction = nn.Linear(self.bertpooler_sec2vec.get_output_dim(),
                                                       self.args.dimentionReductionToThisDim)

    def huggingface_nameloader(self):
        if self.args.bert_name == 'bert-base-uncased':
            self.bert_weight_filepath = 'bert-base-uncased'
        else:
            self.bert_weight_filepath = 'dummy'
            print('Currently not supported', self.args.bert_name)
            exit()

    def forward(self, title_and_desc_concatnated_text):
        mask_sent = get_text_field_mask(title_and_desc_concatnated_text)
        entity_emb = self.word_embedder(title_and_desc_concatnated_text)
        entity_emb = self.word_embedding_dropout(entity_emb)

        if self.args.entityPooling == "CLSLinear":
            entity_emb = entity_emb[:, 0, :]
            entity_emb = self.linear_for_entity_encoding(entity_emb)
        elif self.args.entityPooling == 'CLS':
            entity_emb = entity_emb[:, 0, :]
        else:
            assert self.args.entityPooling == "CLSLinearTanh"
            entity_emb = self.bertpooler_sec2vec(entity_emb, mask_sent)

        if self.args.dimentionReduction:
            return self.linear_for_dimentionReduction(entity_emb)
        else:
            return entity_emb

class Pooler_for_mention(Seq2VecEncoder):
    def __init__(self, args, word_embedder):
        super(Pooler_for_mention, self).__init__()
        self.args = args
        self.huggingface_nameloader()
        self.bertpooler_sec2vec = BertPooler(pretrained_model=self.bert_weight_filepath)
        self.word_embedder = word_embedder
        self.word_embedding_dropout = nn.Dropout(self.args.word_embedding_dropout)

        self.linear_for_mention_encoding = nn.Linear(self.bertpooler_sec2vec.get_output_dim(),
                                                     self.bertpooler_sec2vec.get_output_dim())
        self.linear_for_dimentionReduction = nn.Linear(self.bertpooler_sec2vec.get_output_dim(),
                                                       self.args.dimentionReductionToThisDim)

    def huggingface_nameloader(self):
        if self.args.bert_name == 'bert-base-uncased':
            self.bert_weight_filepath = 'bert-base-uncased'
        else:
            self.bert_weight_filepath = 'dummy'
            print('Currently not supported', self.args.bert_name)
            exit()

    def forward(self, contextualized_mention):
        mask_sent = get_text_field_mask(contextualized_mention)
        mention_emb = self.word_embedder(contextualized_mention)
        mention_emb = self.word_embedding_dropout(mention_emb)

        if self.args.metionPooling == "CLSLinear":
            mention_emb = mention_emb[:, 0, :]
            mention_emb = self.linear_for_mention_encoding(mention_emb)
        elif self.args.metionPooling == 'CLS':
            mention_emb = mention_emb[:, 0, :]
        else:
            assert self.args.metionPooling == "CLSLinearTanh"
            mention_emb = self.bertpooler_sec2vec(mention_emb, mask_sent)

        if self.args.dimentionReduction:
            return self.linear_for_dimentionReduction(mention_emb)
        else:
            return mention_emb

class InKBAllEntitiesEncoder:
    def __init__(self, args, entity_loader_datasetreaderclass, entity_encoder_wrapping_model, vocab):
        self.args = args
        self.entity_loader_datasetreader = entity_loader_datasetreaderclass
        self.sequence_iterator_for_encoding_entities = BasicIterator(batch_size=args.batch_size_for_kb_encoder)
        self.vocab = vocab
        self.entity_encoder_wrapping_model = entity_encoder_wrapping_model
        self.entity_encoder_wrapping_model.eval()
        self.cuda_device = 0

    def encoding_all_entities(self):
        duidx2emb = {}
        ds = self.entity_loader_datasetreader.read('test')
        self.sequence_iterator_for_encoding_entities.index_with(self.vocab)
        entity_generator = self.sequence_iterator_for_encoding_entities(ds, num_epochs=1, shuffle=False)
        entity_generator_tqdm = tqdm(entity_generator, total=self.sequence_iterator_for_encoding_entities.get_num_batches(ds))
        print('======Encoding all entites from title and description=====')
        
        entities_full_path = os.path.join(self.args.entities_path, self.args.entities_filename)
        if self.args.load_entities:
            duidx2emb = pickle_load_object(entities_full_path)
        else:
            with torch.no_grad():
                for batch in entity_generator_tqdm:
                    batch = nn_util.move_to_device(batch, self.cuda_device)
                    duidxs, embs = self._extract_cuidx_and_its_encoded_emb(batch)
                    for duidx, emb in zip(duidxs, embs):
                        duidx2emb.update({int(duidx):emb})
            if self.args.save_entities:
                pickle_save_object(duidx2emb, entities_full_path)

        return duidx2emb

    def tonp(self, tsr):
        return tsr.detach().cpu().numpy()

    def _extract_cuidx_and_its_encoded_emb(self, batch) -> np.ndarray:
        out_dict = self.entity_encoder_wrapping_model(**batch)
        return self.tonp(out_dict['gold_duidx']), self.tonp(out_dict['emb_of_entities_encoded'])

class BiEncoderForOnlyMentionOutput(Model):
    def __init__(self, args,
                 mention_encoder: Seq2VecEncoder,
                 vocab):
        super().__init__(vocab)
        self.args = args
        self.mention_encoder = mention_encoder

    def forward(self, context, gold_title_and_desc_concatenated, gold_duidx, mention_uniq_id,
                gold_and_negatives_title_and_desc_concatenated, labels_for_BCEWithLogitsLoss):
        contextualized_mention = self.mention_encoder(context)
        output = {'mention_uniq_id': mention_uniq_id,
                  'gold_duidx': gold_duidx,
                  'contextualized_mention': contextualized_mention}

        return output
