from tensorflow.keras.models import load_model
import numpy as np
import pickle
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Load Trained Model
model = load_model("lstm_multi_output_model.keras")

# Load Tokenizer
with open("tokenizer.pkl", "rb") as handle:
    tokenizer = pickle.load(handle)

# Define max sequence length (should match training phase)
MAX_LENGTH = 100

# Mapping for question type labels
question_type_dict = {0: "Fill in the Blank", 1: "True/False", 2: "Multiple Choice"}


def predict_question_answer_type(context_text):
    """
    Given a context (paragraph), predict a question, answer, and question type.
    """
    # === Tokenize and Pad the Context === #
    context_seq = tokenizer.texts_to_sequences([context_text])
    context_padded = pad_sequences(
        context_seq, maxlen=MAX_LENGTH, padding="post", truncating="post"
    )

    # === Create a Dummy Input for the Question Decoder === #
    # Since our model is seq2seq, we initialize an empty decoder input
    decoder_input_seq = np.zeros((1, MAX_LENGTH))

    # === Predict Outputs === #
    predicted_question_seq, predicted_answer_seq, predicted_type = model.predict(
        [context_padded, decoder_input_seq]
    )

    # === Convert Predicted Sequences to Text === #
    index_to_word = {v: k for k, v in tokenizer.word_index.items()}

    # Convert predicted question
    predicted_question = " ".join(
        [
            index_to_word[idx]
            for idx in np.argmax(predicted_question_seq, axis=2)[0]
            if idx in index_to_word
        ]
    )

    # Convert predicted answer
    predicted_answer = " ".join(
        [
            index_to_word[idx]
            for idx in np.argmax(predicted_answer_seq, axis=2)[0]
            if idx in index_to_word
        ]
    )

    # Convert predicted question type (numerical label → text)
    predicted_question_type = question_type_dict[np.argmax(predicted_type)]

    return predicted_question, predicted_answer, predicted_question_type


# Sample Test Context
context_example = "Ki Hajar Dewantara adalah pelopor pendidikan di Indonesia dan pendiri Taman Siswa. Ia dikenal dengan semboyannya 'Ing Ngarsa Sung Tuladha, Ing Madya Mangun Karsa, Tut Wuri Handayani', yang menekankan peran guru dalam pendidikan."

# Run the prediction
predicted_question, predicted_answer, predicted_question_type = (
    predict_question_answer_type(context_example)
)

# Print the Results
print(f"🔹 Predicted Question: {predicted_question}")
print(f"🔹 Predicted Answer: {predicted_answer}")
print(f"🔹 Predicted Question Type: {predicted_question_type}")