import numpy as np

import tensorflow as tf

from typing import List

from nltk.translate.bleu_score import sentence_bleu

from modules.chatbot.preprocessor import preprocess

class Inferencer:

    def __init__(

        self,

        medical_qa_gpt_model: tf.keras.Model,

        bert_tokenizer: tf.keras.preprocessing.text.Tokenizer,

        gpt_tokenizer: tf.keras.preprocessing.text.Tokenizer,

        question_extractor_model: tf.keras.Model,

        df_qa: pd.DataFrame,

        answer_index: faiss.IndexFlatIP,

        answer_len: int,

    ) -> None:

        """

        Initialize Inferencer with necessary components.

        Args:

            medical_qa_gpt_model (tf.keras.Model): Medical Q&A GPT model.

            bert_tokenizer (tf.keras.preprocessing.text.Tokenizer): BERT tokenizer.

            gpt_tokenizer (tf.keras.preprocessing.text.Tokenizer): GPT tokenizer.

            question_extractor_model (tf.keras.Model): Question extractor model.

            df_qa (pd.DataFrame): DataFrame containing Q&A pairs.

            answer_index (faiss.IndexFlatIP): FAISS index for answers.

            answer_len (int): Length of the answer.

        """

        self.biobert_tokenizer = bert_tokenizer

        self.question_extractor_model = question_extractor_model

        self.answer_index = answer_index

        self.gpt_tokenizer = gpt_tokenizer

        self.medical_qa_gpt_model = medical_qa_gpt_model

        self.df_qa = df_qa

        self.answer_len = answer_len

    def get_gpt_inference_data(

        self, question: str, question_embedding: np.ndarray

    ) -> List[int]:

        """

        Get GPT inference data.

        Args:

            question (str): Input question.

            question_embedding (np.ndarray): Embedding of the question.

        Returns:

            List[int]: GPT inference data.

        """

        topk = 20

        scores, indices = self.answer_index.search(

            question_embedding.astype("float32"), topk

        )

        q_sub = self.df_qa.iloc[indices.reshape(20)]

        line = "`QUESTION: %s `ANSWER: " % (question)

        encoded_len = len(self.gpt_tokenizer.encode(line))

        for i in q_sub.iterrows():

            line = (

                "`QUESTION: %s `ANSWER: %s " % (i[1]["question"], i[1]["answer"]) + line

            )

            line = line.replace("\n", "")

            encoded_len = len(self.gpt_tokenizer.encode(line))

            if encoded_len >= 1024:

                break

        return self.gpt_tokenizer.encode(line)[-1024:]

    def get_gpt_answer(self, question: str, answer_len: int) -> str:

        """

        Get GPT answer.

        Args:

            question (str): Input question.

            answer_len (int): Length of the answer.

        Returns:

            str: GPT generated answer.

        """

        preprocessed_question = preprocess(question)

        truncated_question = (

            " ".join(preprocessed_question.split(" ")[:500])

            if len(preprocessed_question.split(" ")) > 500

            else preprocessed_question

        )

        encoded_question = self.biobert_tokenizer.encode(truncated_question)

        padded_question = tf.keras.preprocessing.sequence.pad_sequences(

            [encoded_question], maxlen=512, padding="post"

        )

        question_mask = np.where(padded_question != 0, 1, 0)

        embeddings = self.question_extractor_model(

            {"question": padded_question, "question_mask": question_mask}

        )

        gpt_input = self.get_gpt_inference_data(truncated_question, embeddings.numpy())

        mask_start = len(gpt_input) - list(gpt_input[::-1]).index(4600) + 1

        input = gpt_input[: mask_start + 1]

        if len(input) > (1024 - answer_len):

            input = input[-(1024 - answer_len) :]

        gpt2_output = self.gpt_tokenizer.decode(

            self.medical_qa_gpt_model.generate(

                input_ids=tf.constant([np.array(input)]),

                max_length=1024,

                temperature=0.7,

            )[0]

        )

        answer = gpt2_output.rindex("`ANSWER: ")

        return gpt2_output[answer + len("`ANSWER: ") :]

    def inf_func(self, question: str) -> str:

        """

        Run inference for the given question.

        Args:

            question (str): Input question.

        Returns:

            str: Generated answer.

        """

        answer_len = self.answer_len

        return self.get_gpt_answer(question, answer_len)

    def eval_func(self, question: str, answer: str) -> float:

        """

        Evaluate generated answer against ground truth.

        Args:

            question (str): Input question.

            answer (str): Generated answer.

        Returns:

            float: BLEU score.

        """

        answer_len = 20

        generated_answer = self.get_gpt_answer(question, answer_len)

        reference = [answer.split(" ")]

        candidate = generated_answer.split(" ")

        score = sentence_bleu(reference, candidate)

        return score

    def run(self, question: str, isEval: bool) -> str:

        """

        Run inference for the given question.

        Args:

            question (str): Input question.

            isEval (bool): Whether to evaluate or not.

        Returns:

            str: Generated answer.

        """

        answer = self.inf_func(question)

        if isEval:

            bleu_score = self.eval_func(question, answer)

            print(f"The sentence_bleu score is {bleu_score}")

        return answer