import sys
import math
import numpy as np
import torch
from scipy import stats
from torch.nn.functional import softmax
import keras
if(keras.__version__ == '2.7.0'):
    from keras.preprocessing.sequence import pad_sequences
else:
    from keras.utils import pad_sequences
from torch.utils.data import TensorDataset, SequentialSampler, DataLoader
from transformers import PreTrainedTokenizer
from typing import Tuple
import pandas as pd
import regex as re
from copy import copy
from BiasDetection.metrics.LogProbability.LogProbability import LogProbability
from BiasDetection.metrics.LogProbability.crows.crows import CrowSPairsRunner
class LogProbabilityGender(LogProbability):
    def __init__(self, model, tokenizer, device, model_type, model_class, mask_token='[MASK]', dataset='crows'):
        super().__init__(model, tokenizer, device, model_type, model_class, mask_token, dataset)
        self.bec_pro_path = sys.path[1]+'data/BEC-Pro/Bec-Pro_EN.tsv'
        self.dataset = dataset
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.model_type = model_type
        self.mask_token = mask_token
    def statistics(self, group1, group2):
        """take 2 groups of paired samples and compute either a paired samples t-test or
        a Wilcoxon signed rank test
        prints out a description of the two groups as well as the statistic and p value of the test"""
        assert len(group1) == len(group2), "The two groups do not have the same length"

        print('Group 1:')
        print(group1.describe())
        print('Group 2:')
        print(group2.describe())

        dif = group1.sub(group2, fill_value=0)

        SW_stat, SW_p = stats.shapiro(dif)
        print(SW_stat, SW_p)

        if SW_p >= 0.05:
            print('T-Test:')
            statistic, p = stats.ttest_rel(group1, group2)
        else:
            print('Wilcoxon Test:')
            statistic, p = stats.wilcoxon(group1, group2)

        print('Statistic: {}, p: {}'.format(statistic, p))

        effect_size = statistic / np.sqrt(len(group1))
        print('effect size r: {}'.format(effect_size))

        return

    def attention_mask_creator(self, input_ids):
        """Provide the attention mask list of lists: 0 only for [PAD] tokens (index 0)
        Returns torch tensor"""
        attention_masks = []
        for sent in input_ids:
            segments_ids = [int(t > 0) for t in sent]
            attention_masks.append(segments_ids)
        return torch.tensor(attention_masks)

    def tokenize_to_id(self, sentences, tokenizer):
        """Tokenize all of the sentences and map the tokens to their word IDs."""
        input_ids = []

        # For every sentence...
        for sent in sentences:
            # `encode` will:
            #   (1) Tokenize the sentence.
            #   (2) Prepend the `[CLS]` token to the start.
            #   (3) Append the `[SEP]` token to the end.
            #   (4) Map tokens to their IDs.
            encoded_sent = tokenizer.encode(
                sent,  # Sentence to encode.
                add_special_tokens=True,  # Add '[CLS]' and '[SEP]'

                # This function also supports truncation and conversion
                # to pytorch tensors, but we need to do padding, so we
                # can't use these features :( .
                # max_length = 128,          # Truncate all sentences.
                # return_tensors = 'pt',     # Return pytorch tensors.
            )

            # Add the encoded sentence to the list.
            input_ids.append(encoded_sent)
        return input_ids

    def input_pipeline(self, sequence, tokenizer, MAX_LEN):
        """function to tokenize, pad and create attention masks"""
        input_ids = self.tokenize_to_id(sequence, tokenizer)
        input_ids = pad_sequences(input_ids, maxlen=MAX_LEN,
                                dtype="long",
                                value=tokenizer.mask_token_id,
                                truncating="post", padding="post")
        input_ids = torch.tensor(input_ids)

        attention_masks = self.attention_mask_creator(input_ids)

        return input_ids, attention_masks


    def prob_with_prior(self, pred_TM, pred_TAM, input_ids_TAM, original_ids, tokenizer):
        pred_TM = pred_TM.cpu()
        pred_TAM = pred_TAM.cpu()
        input_ids_TAM = input_ids_TAM.cpu()

        probs = []
        for doc_idx, id_list in enumerate(input_ids_TAM):
            # see where the masks were placed in this sentence
            mask_indices = np.where(id_list == tokenizer.mask_token_id)[0]
            # now get the probability of the target word:
            # first get id of target word
            target_id = original_ids[doc_idx][mask_indices[0]]
            # get its probability with unmasked profession
            target_prob = pred_TM[doc_idx][mask_indices[0]][target_id].item()
            # get its prior probability (masked profession)
            prior = pred_TAM[doc_idx][mask_indices[0]][target_id].item()

            probs.append(np.log(target_prob / prior))

        return probs


    def LogProbabilityBec(self, model, tokenizer, device):
        eval_df = pd.read_csv(self.bec_pro_path, sep='\t')
        """takes professional sentences as DF, a tokenizer & a BERTformaskedLM model
        and predicts the associations"""

        # as max_len get the smallest power of 2 greater or equal to the max sentence lenght
        max_len = max([len(sent.split()) for sent in eval_df.Sent_TM])
        pos = math.ceil(math.log2(max_len))
        max_len_eval = int(math.pow(2, pos))

        print('max_len evaluation: {}'.format(max_len_eval))

        # create BERT-ready inputs: target masked, target and attribute masked,
        # and the tokenized original inputs to recover the original target word
        eval_tokens_TM, eval_attentions_TM = self.input_pipeline(eval_df.Sent_TM,
                                                            tokenizer,
                                                            max_len_eval)
        eval_tokens_TAM, eval_attentions_TAM = self.input_pipeline(eval_df.Sent_TAM,
                                                            tokenizer,
                                                            max_len_eval)
        eval_tokens, _ = self.input_pipeline(eval_df.Sentence, tokenizer, max_len_eval)

        # check that lengths match before going further
        assert eval_tokens_TM.shape == eval_attentions_TM.shape
        assert eval_tokens_TAM.shape == eval_attentions_TAM.shape

        # make a Evaluation Dataloader
        eval_batch = 20
        eval_data = TensorDataset(eval_tokens_TM, eval_attentions_TM,
                                eval_tokens_TAM, eval_attentions_TAM,
                                eval_tokens)
        eval_sampler = SequentialSampler(eval_data)
        eval_dataloader = DataLoader(eval_data, batch_size=eval_batch,
                                    sampler=eval_sampler)

        # put everything to GPU (if it is available)
        # eval_tokens_TM = eval_tokens_TM.to(device)
        # eval_attentions_TM = eval_attentions_TM.to(device)
        # eval_tokens_TAM = eval_tokens_TAM.to(device)
        # eval_attentions_TAM = eval_attentions_TAM.to(device)
        model.to(device)

        # put model in evaluation mode & start predicting
        model.eval()
        associations_all = []
        for step, batch in enumerate(eval_dataloader):
            b_input_TM = batch[0].to(device)
            b_att_TM = batch[1].to(device)
            b_input_TAM = batch[2].to(device)
            b_att_TAM = batch[3].to(device)

            with torch.no_grad():
                outputs_TM = model(b_input_TM,
                                attention_mask=b_att_TM)
                outputs_TAM = model(b_input_TAM,
                                    attention_mask=b_att_TAM)
                predictions_TM = softmax(outputs_TM[0], dim=2)
                predictions_TAM = softmax(outputs_TAM[0], dim=2)

            assert predictions_TM.shape == predictions_TAM.shape

            # calculate associations
            associations = self.prob_with_prior(predictions_TM,
                                        predictions_TAM,
                                        b_input_TAM,
                                        batch[4],  # normal inputs
                                        tokenizer)

            associations_all += associations

        return associations_all
    def LogProbabilityCrows(self, model, tokenizer, device, model_type):
        is_generative = True
        if(model_type == 'masked'):
            is_generative = False
        runner = CrowSPairsRunner(
            model=model,
            tokenizer=tokenizer,
            input_file=self.crows_path,
            bias_type='gender',
            is_generative=is_generative  # Affects model scoring.
        )
        results = runner()

        print(f"Metric: {results}")
        return results
    
    def evaluate(self):
        if('crows' in self.dataset):
            return self.LogProbabilityCrows(self.model, self.tokenizer, self.device, self.model_type)
        if('bec' in self.dataset):
            return self.LogProbabilityBec(self.model, self.tokenizer, self.device, self.model_type)
