import spaces

import torch
import gradio as gr
from transformers import AutoTokenizer

from custom_model import MCQBiEncoder
from predict import predict

MODEL_PATH = "custom_model.pt"
MAX_LENGTH = 256
device = "cuda"

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

model = MCQBiEncoder(
    vocab_size=tokenizer.vocab_size,
    dropout=0.1,
    temperature=0.07
)

state_dict = torch.load(
    MODEL_PATH,
    map_location="cpu"
)

model.load_state_dict(state_dict)
model.to(device)
model.eval()

@spaces.GPU
def solve_mcq(
    question,
    option_a,
    option_b,
    option_c,
    option_d,
    option_e
):

    options = [
        option_a,
        option_b,
        option_c,
        option_d,
        option_e
    ]

    question_tokens = tokenizer(
        question,
        padding="max_length",
        truncation=True,
        max_length=MAX_LENGTH
    )

    option_tokens = [
        tokenizer(
            option,
            padding="max_length",
            truncation=True,
            max_length=MAX_LENGTH
        )
        for option in options
    ]

    q_ids = torch.tensor(
        [question_tokens["input_ids"]],
        dtype=torch.long,
        device=device
    )

    q_mask = torch.tensor(
        [question_tokens["attention_mask"]],
        dtype=torch.long,
        device=device
    )

    opt_ids = torch.tensor(
        [[option["input_ids"] for option in option_tokens]],
        dtype=torch.long,
        device=device
    )

    opt_mask = torch.tensor(
        [[option["attention_mask"] for option in option_tokens]],
        dtype=torch.long,
        device=device
    )

    scores, top3 = predict(
        model,
        q_ids,
        q_mask,
        opt_ids,
        opt_mask,
        device
    )

    answer_labels = [
        "A",
        "B",
        "C",
        "D",
        "E"
    ]

    top3 = top3[0].tolist()

    results = []

    for index in top3:

        results.append(
            f"{answer_labels[index]}. {options[index]}"
        )

    return "\n\n".join(results)


demo = gr.Interface(
    fn=solve_mcq,

    inputs=[
        gr.Textbox(
            label="Question",
            placeholder="Enter your question..."
        ),

        gr.Textbox(
            label="Option A"
        ),

        gr.Textbox(
            label="Option B"
        ),

        gr.Textbox(
            label="Option C"
        ),

        gr.Textbox(
            label="Option D"
        ),

        gr.Textbox(
            label="Option E"
        ),
    ],

    outputs=gr.Textbox(
        label="Top 3 Answers"
    ),

    title="MCQ Solver",

    description=(
        "Enter a question and five options. "
        "The model ranks the three most likely answers."
    )
)


if __name__ == "__main__":
    demo.launch()