import json
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model  # type: ignore
from tensorflow.keras.preprocessing.sequence import pad_sequences  # type: ignore
from tensorflow.keras.preprocessing.text import tokenizer_from_json  # type: ignore
import re


class AnswerGenerationRepository:
    MODEL_PATH = "app/lstm_model/question_generation/qa_lstm_model_final.h5"
    TOKENIZER_PATH = "app/lstm_model/question_generation/qa_tokenizers.json"

    def __init__(self):
        with open(self.TOKENIZER_PATH, "r") as f:
            tokenizer_data = json.load(f)

        self.tokenizer = tokenizer_from_json(tokenizer_data["word_tokenizer"])
        self.ner_tokenizer = tokenizer_from_json(tokenizer_data["ner_tokenizer"])
        self.srl_tokenizer = tokenizer_from_json(tokenizer_data["srl_tokenizer"])
        self.answer_tokenizer = tokenizer_from_json(tokenizer_data["answer_tokenizer"])
        self.q_type_tokenizer = tokenizer_from_json(tokenizer_data["q_type_tokenizer"])

        self.max_context_len = tokenizer_data["max_context_len"]
        self.max_question_len = tokenizer_data["max_question_len"]
        self.max_token_len = tokenizer_data["max_token_len"]
        self.q_type_vocab_size = len(self.q_type_tokenizer.word_index) + 1

        self.model = load_model(self.MODEL_PATH)

    def preprocess_text(self, text):
        text = text.lower()
        text = re.sub(r"\s+", " ", text).strip()
        return text

    def predict_answer(self, context, question, tokens, ner, srl, q_type):
        context_seq = self.tokenizer.texts_to_sequences([self.preprocess_text(context)])
        question_seq = self.tokenizer.texts_to_sequences(
            [self.preprocess_text(question)]
        )
        token_seq = [self.tokenizer.texts_to_sequences([" ".join(tokens)])[0]]
        ner_seq = [self.ner_tokenizer.texts_to_sequences([" ".join(ner)])[0]]
        srl_seq = [self.srl_tokenizer.texts_to_sequences([" ".join(srl)])[0]]

        q_type_idx = self.q_type_tokenizer.word_index.get(q_type, 0)
        q_type_cat = tf.keras.utils.to_categorical(
            [q_type_idx], num_classes=self.q_type_vocab_size
        )

        context_pad = pad_sequences(
            context_seq, maxlen=self.max_context_len, padding="post"
        )
        question_pad = pad_sequences(
            question_seq, maxlen=self.max_question_len, padding="post"
        )
        token_pad = pad_sequences(token_seq, maxlen=self.max_token_len, padding="post")
        ner_pad = pad_sequences(ner_seq, maxlen=self.max_token_len, padding="post")
        srl_pad = pad_sequences(srl_seq, maxlen=self.max_token_len, padding="post")

        prediction = self.model.predict(
            [context_pad, question_pad, token_pad, ner_pad, srl_pad, q_type_cat],
            verbose=0,
        )
        answer_idx = np.argmax(prediction[0])

        for word, idx in self.answer_tokenizer.word_index.items():
            if idx == answer_idx:
                return word

        return "Unknown"
