import itertools
import os
from typing import List

import numpy as np
import torch
from torch.nn.functional import one_hot
from torch.nn.utils.rnn import pad_sequence
from transformers import BertForMaskedLM, BertTokenizer
from transformers import TransfoXLLMHeadModel, TransfoXLTokenizer

from candidate import Candidate as CandidateUnit


class ContextProbability(object):
    def __init__(self, **kwargs):
        self.device = kwargs['device'] if 'device' in kwargs else 'cpu'
        print('Using ' + str(self.device))

    def _softmax(self, values):
        values = np.asarray(values).reshape(-1)
        max_values = np.max(values)
        values = values - max_values
        exp_values = np.exp(values)
        softmax_scores = exp_values / exp_values.sum()
        return softmax_scores.tolist()


class ContextProbabilityBERTLM(ContextProbability):

    def __init__(self, pretrained_model_path="", **kwargs):
        super().__init__(**kwargs)
        self.__mask_token = "[MASK]"
        self.__cls_token = "[CLS]"
        self.__sep_token = "[SEP]"
        self.__model_class = BertForMaskedLM
        self.__tokenizer_class = BertTokenizer
        self.__pretrained_type = 'bert-base-uncased'
        self.model, self.tokenizer = self._load_bert_model(pretrained_model_path)
        self.vocab = self._get_bert_vocab()

    def _load_bert_model(self, path):
        try:
            assert os.path.exists(path) == True
            model = self.__model_class.from_pretrained(path)
            tokenizer = self.__tokenizer_class.from_pretrained(path)
        except:
            print("Unable to load model and tokenizer data from local path: \
                 {}; Loading from huggingface site".format(path))
            model = self.__model_class.from_pretrained(self.__pretrained_type)
            tokenizer = self.__tokenizer_class.from_pretrained(self.__pretrained_type)
        model = model.to(self.device)
        return model, tokenizer

    def _get_bert_vocab(self):
        if self.tokenizer == None:
            self.model, self.tokenizer = self._load_bert_model()
        return self.tokenizer.vocab

    def _get_masked_sents(self, tokenized_sentences, indices, add_terminal_tokens=False):
        assert len(tokenized_sentences) == len(indices)
        masked_sents = []
        for (tsent, ind) in zip(tokenized_sentences, indices):
            tsent[ind] = self.__mask_token
            msent = "{} {} {}".format(self.__cls_token, " ".join(tsent), self.__sep_token) \
                if add_terminal_tokens else " ".join(tsent)
            masked_sents.append(msent)
        return masked_sents

    def __addto_forward_pass_stacklist(self, wtsent, masked_sents, tsent_seq_len, tsent_mask_index, candidate_map, i,
                                       candidate, vocab_search_candidate):
        potential_msent = " ".join(wtsent)
        if potential_msent in masked_sents:
            candidate_map[i][candidate].append(
                (masked_sents.index(potential_msent), self.vocab[vocab_search_candidate]))
        else:
            masked_sents.append(potential_msent)
            dummy1 = self.tokenizer.tokenize(potential_msent)
            tsent_mask_index.append(
                dummy1.index(self.__mask_token))  # returns first index of self.__mask_token if multiple exists
            tsent_seq_len.append(len(dummy1))
            candidate_map[i][candidate].append((len(masked_sents) - 1, self.vocab[vocab_search_candidate]))
        return masked_sents, tsent_seq_len, tsent_mask_index, candidate_map

    def _get_bert_probabilities(self, tokenized_sentences, indices, candidates=[],
                                max_seq_len=-1) -> "List[List[CandidateUnit]]":

        for i, wtsent in enumerate(tokenized_sentences):
            tokenized_sentences[i] = [self.__cls_token] + wtsent + [self.__sep_token]

        masked_sents, tsent_seq_len, tsent_mask_index = [], [], []
        if len(candidates) > 0:
            assert len(candidates) == len(tokenized_sentences) == len(indices)
            candidate_map = {}
            for i, (wtsent, wtind, row) in enumerate(zip(tokenized_sentences, indices, candidates)):
                candidate_map[i] = {}
                for candidate in row:
                    candidate_map[i][candidate] = []
                    if candidate in self.vocab:
                        wtsent_copy = list(wtsent)
                        wtsent_copy[wtind] = self.__mask_token
                        masked_sents, tsent_seq_len, tsent_mask_index, candidate_map = \
                            self.__addto_forward_pass_stacklist(wtsent_copy, masked_sents, tsent_seq_len,
                                                                tsent_mask_index, candidate_map, i, candidate,
                                                                candidate)
                    else:
                        # if broken_candidates = [c1, c2, c3, c4], we have four potential_msents with different masking scheme
                        broken_candidates = self.tokenizer.tokenize(candidate)
                        n_broken = len(broken_candidates)
                        replace_with = [self.__mask_token] * n_broken
                        wtsent_copy = list(wtsent)
                        wtsent_copy[wtind] = replace_with
                        wtsent_copy = list(itertools.chain.from_iterable(
                            itertools.repeat(x, 1) if isinstance(x, str) else x for x in wtsent_copy))
                        masked_sents, tsent_seq_len, tsent_mask_index, candidate_map = \
                            self.__addto_forward_pass_stacklist(wtsent_copy, masked_sents, tsent_seq_len,
                                                                tsent_mask_index, candidate_map, i, candidate,
                                                                broken_candidates[0])
                        for temp_i, br_candidate in enumerate(broken_candidates[:-1]):
                            replace_with[temp_i] = br_candidate
                            wtsent_copy = list(wtsent)
                            wtsent_copy[wtind] = replace_with
                            wtsent_copy = list(itertools.chain.from_iterable(
                                itertools.repeat(x, 1) if isinstance(x, str) else x for x in wtsent_copy))
                            masked_sents, tsent_seq_len, tsent_mask_index, candidate_map = \
                                self.__addto_forward_pass_stacklist(wtsent_copy, masked_sents, tsent_seq_len,
                                                                    tsent_mask_index, candidate_map, i, candidate,
                                                                    broken_candidates[temp_i + 1])
            return_selected_candidates = True
            # print(candidate_map)
            # print(masked_sents)
        else:
            assert len(tokenized_sentences) == len(indices)
            masked_sents = self._get_masked_sents(tokenized_sentences, indices, add_terminal_tokens=True)
            for i, msent in enumerate(masked_sents):
                tsent = self.tokenizer.tokenize(msent)
                tsent_seq_len.append(len(tsent))
                tsent_mask_index.append(tsent.index(self.__mask_token))
            return_selected_candidates = False

        nsents = len(masked_sents)

        BATCH_SIZE, batch_start = 128, 0
        # obtain softmax scores as list; if converted to numpy has shape [nsents,bert_vocab_size]
        softmax_scores = []
        while (batch_start < nsents):

            batch_masked_sents = masked_sents[batch_start:min(batch_start + BATCH_SIZE, nsents)]
            batch_tsent_seq_len = tsent_seq_len[batch_start:min(batch_start + BATCH_SIZE, nsents)]
            batch_tsent_mask_index = tsent_mask_index[batch_start:min(batch_start + BATCH_SIZE, nsents)]

            # max_seq_len must be at least 1 greater than until MASK index length!!
            if max_seq_len != -1 and len([*filter(lambda x: x <= 0, max_seq_len - batch_tsent_seq_len - 2)]) == 0:
                MAX_SEQUENCE_LEN = max_seq_len
            else:
                MAX_SEQUENCE_LEN = np.max(batch_tsent_seq_len)

            batch_encoded_sents = []
            for msent in batch_masked_sents:
                tsent = self.tokenizer.tokenize(msent)
                if len(tsent) > MAX_SEQUENCE_LEN:
                    tsent = tsent[:MAX_SEQUENCE_LEN]
                    tsent[-1] = ["[SEP]"]
                tsent += ["[PAD]"] * (MAX_SEQUENCE_LEN - len(tsent))
                idsent = self.tokenizer.convert_tokens_to_ids(tsent)
                batch_encoded_sents.append(idsent)
            batch_input_ids = torch.tensor(batch_encoded_sents)
            batch_input_ids = batch_input_ids.to(self.device)

            with torch.no_grad():
                batch_final_layer_logits = self.model(batch_input_ids)[0]  # [BATCH_SIZE,MAX_SEQUENCE_LEN,32k]

            batch_final_layer_logits = batch_final_layer_logits.to('cpu')
            for _final_layer_logits, _tsent_mask_index in zip(batch_final_layer_logits, batch_tsent_mask_index):
                curr_softmax_scores = self._softmax(_final_layer_logits[_tsent_mask_index, :].reshape(-1))
                softmax_scores.append(curr_softmax_scores)

            batch_start += BATCH_SIZE

        if return_selected_candidates:
            return_candidate_instances = []
            for i in candidate_map.keys():
                curr_candidate_instances = []
                for candidate in candidate_map[i].keys():
                    # can have a list of posibilities as [(x1,y1),...,(xn,yn)]
                    #   due to split into sub word units
                    prod = 1
                    for (ind1, ind2) in candidate_map[i][candidate]:
                        prod *= softmax_scores[ind1][ind2]
                    curr_candidate_instances.append(CandidateUnit(candidate, prod))
                return_candidate_instances.append(curr_candidate_instances)
        else:
            return_candidate_instances = []
            for i in range(0, nsents):
                curr_candidate_instances = []
                for (candidate, score) in zip([*self.vocab.keys()], softmax_scores[i]):
                    curr_candidate_instances.append(CandidateUnit(candidate, score))
                return_candidate_instances.append(curr_candidate_instances)

        # sort
        for i, candidate_list in enumerate(return_candidate_instances):
            return_candidate_instances[i] = sorted(candidate_list, key=lambda unit_: unit_.prob, reverse=True)

        return return_candidate_instances


class SentenceProbability:
    def __init__(self, device='cpu'):
        tokenizer = TransfoXLTokenizer.from_pretrained('transfo-xl-wt103')
        model = TransfoXLLMHeadModel.from_pretrained('transfo-xl-wt103')
        model = model.to(device)

        self.tokenizer = tokenizer
        self.model = model.eval()
        self.device = device

        self.NUM_CLASSES = 267735

    def get_sentence_probabilities(self, sentences: List[List[str]]) -> torch.Tensor:
        """
        Given a list of list of words (list of sentences), this function computes the
        probability of the complete sentence according to the TransformerXL model.

        Returns a 1-D tensor with probability values for each sentence.
        """
        sentences = sorted(sentences, key=lambda sentence: len(sentence))

        input_ids = self.tokenizer.encode_sents(sentences)
        input_mask = [torch.ones_like(input_tensor) for input_tensor in input_ids]

        padded_input_tensor = pad_sequence(input_ids, batch_first=True).to(self.device)
        padded_mask_tensor = pad_sequence(input_mask, batch_first=True).float().to(self.device)

        with torch.no_grad():
            prediction_scores, _ = self.model(padded_input_tensor)[:2]

        expanded_input_tensor = one_hot(padded_input_tensor, num_classes=self.NUM_CLASSES).float()
        word_wise_probabilities = torch.mul(expanded_input_tensor, prediction_scores).sum(dim=2)

        masked_word_wise_probabilities = torch.mul(word_wise_probabilities, padded_mask_tensor)
        sentence_wise_probabilities = masked_word_wise_probabilities.sum(dim=1)

        return sentence_wise_probabilities.to('cpu')

    def get_candidate_probabilities(self,
                                    sentences: List[List[str]],
                                    indices: List[int],
                                    candidates: List[List[str]]) -> torch.Tensor:
        """
        Given sentences, indices for each sentence, and potential candidates
        for those indices, this function returns a 2D tensor containing
        probabilities for each sentence and candidate pair in the same order.

        Assumption - All words have equal number of candidates.
        """
        all_probabilities = []
        for sentence, index, candidate_list in zip(sentences, indices, candidates):
            candidate_sentence_list = []
            for candidate in candidate_list:
                new_sentence = sentence.copy()
                new_sentence[index] = candidate
                candidate_sentence_list.append(new_sentence)

            all_probabilities.append(self.get_sentence_probabilities(candidate_sentence_list))

        return torch.stack(all_probabilities)
