import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

from interfaces.QuestionAnswerOperation import QuestionAnswerOperation
from tasks.TaskTypes import TaskType
import numpy as np

""" The following transformation augments question answering data by generating paraphrases of questions. 
It relies on a question-question pair fine-tuned T5 model which has been taken from https://huggingface.co/ramsrigouthamg/t5_paraphraser. 
The test cases were generated by the commented script in redundant_context_for_qa. 

"""


class QuoraT5QaPairGenerator(QuestionAnswerOperation):
    tasks = [
        TaskType.QUESTION_ANSWERING,
        TaskType.QUESTION_GENERATION
    ]
    languages = ["en"]
    heavy = True

    def __init__(self, seed=0, model_name="ramsrigouthamg/t5_paraphraser", max_len=256, num_return_sequences=3, top_k=40):
        super().__init__(seed)
        self.model = T5ForConditionalGeneration.from_pretrained(model_name)
        self.tokenizer = T5Tokenizer.from_pretrained(model_name)
        self.max_len = max_len
        self.num_return_sequences = num_return_sequences
        self.top_k = top_k
        torch.manual_seed(self.seed)
        np.random.seed(self.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(self.seed)

    def generate(self, context: str, question: str, answers: [str]):

        original = question
        text = "paraphrase: " + original + "</s>"
        encoding = self.tokenizer.encode_plus(text, padding="longest", return_tensors="pt")
        input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
        beam_outputs = self.model.generate(
            input_ids=input_ids, attention_mask=attention_masks,
            do_sample=True,
            max_length=self.max_len,
            top_k=self.top_k,
            top_p=0.98,
            early_stopping=True,
            num_return_sequences=self.num_return_sequences
        )
        unique_question_paraprases = []
        for beam_output in beam_outputs:
            paraphrase = self.tokenizer.decode(beam_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
            if paraphrase.strip() and paraphrase.lower() != original.lower() and paraphrase not in unique_question_paraprases:
                unique_question_paraprases.append(paraphrase)

        paraphrased_questions = unique_question_paraprases if len(unique_question_paraprases) > 0 else [original]
        print(paraphrased_questions)
        return [(context, p, answers) for p in paraphrased_questions]
