"""Smart MCQ Solver - Gradio demo.

Loads the fine-tuned DeBERTa-v3-large multiple-choice model from the Hub and
ranks five candidate answers for a question.
"""
import os
import numpy as np
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForMultipleChoice

# Available only on ZeroGPU hardware. On CPU-basic the import fails and we stay on CPU,
# so this one file runs unchanged on either.
try:
    import spaces
    HAS_ZEROGPU = True
except ImportError:
    HAS_ZEROGPU = False

MODEL_ID = os.environ.get("MODEL_ID", "SriragData/smart-mcq-deberta-v3-large")
MAX_LEN = 256
LAB = list("ABCDE")

torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2))

print(f"loading {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# fp32 is forced explicitly on purpose: transformers honours the dtype recorded in
# config.json, and an fp16 DeBERTa-v3 collapses to a constant output with no error
# raised. The kwarg was renamed (torch_dtype -> dtype) between versions, so try both
# rather than pinning ourselves to one.
try:
    model = AutoModelForMultipleChoice.from_pretrained(MODEL_ID, dtype=torch.float32)
except TypeError:
    model = AutoModelForMultipleChoice.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
model = model.eval()
assert next(model.parameters()).dtype == torch.float32, "model did not load in fp32"
print("loaded. params:", sum(p.numel() for p in model.parameters()))


TEMPLATES = [
    "Pick the best possible answer:", "Select the most accurate option:",
    "Identify the correct statement:", "Determine the correct option:",
    "Choose the correct answer:", "Which of the following is correct?",
]


def clean(prompt: str) -> str:
    """Strip the boilerplate wrappers the training data was generated with."""
    for t in TEMPLATES:
        prompt = prompt.replace(t, "")
    return prompt.strip()


@torch.no_grad()
def solve(question, a, b, c, d, e):
    options = [a, b, c, d, e]
    if not question or not question.strip():
        return "Enter a question.", None
    if any(not str(o).strip() for o in options):
        return "All five options are required.", None

    # On ZeroGPU, CUDA only exists inside the decorated call, so resolve the device here
    # rather than at import time.
    device = "cuda" if torch.cuda.is_available() else "cpu"
    m = model.to(device)

    q = clean(question)
    enc = tokenizer([q] * 5, [str(o) for o in options],
                    truncation=True, max_length=MAX_LEN,
                    padding=True, return_tensors="pt")
    logits = m(input_ids=enc["input_ids"].unsqueeze(0).to(device),
               attention_mask=enc["attention_mask"].unsqueeze(0).to(device)).logits[0]
    probs = torch.softmax(logits.float().cpu(), dim=-1).numpy()

    order = np.argsort(-probs)
    top3 = " ".join(LAB[i] for i in order[:3])
    verdict = (f"### {LAB[order[0]]} — {options[order[0]]}\n\n"
               f"**MAP@3 submission format:** `{top3}`")

    rows = [[LAB[i], str(options[i])[:200], round(float(probs[i]), 4),
             int(np.where(order == i)[0][0]) + 1] for i in range(5)]
    rows.sort(key=lambda r: r[3])
    return verdict, rows


if HAS_ZEROGPU:
    # Requests a GPU slice for the duration of the call. Applied programmatically so the
    # same file still imports on CPU-basic hardware, where `spaces` does not exist.
    solve = spaces.GPU(duration=60)(solve)


EXAMPLES = [
    ["Which philosopher argued that existence precedes essence?",
     "Immanuel Kant", "Jean-Paul Sartre", "David Hume", "Rene Descartes", "Baruch Spinoza"],
    ["What is the primary function of mitochondria in a eukaryotic cell?",
     "Protein synthesis", "Storage of genetic material",
     "Production of ATP through oxidative phosphorylation",
     "Breakdown of cellular waste", "Regulation of cell division"],
    ["According to the second law of thermodynamics, what happens to the entropy "
     "of an isolated system over time?",
     "It decreases to zero", "It remains exactly constant",
     "It never decreases", "It oscillates periodically", "It becomes negative"],
]

CSS = """
.gradio-container {max-width: 1000px !important}
footer {visibility: hidden}
"""

with gr.Blocks(title="Smart MCQ Solver", css=CSS, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# Smart MCQ Solver\n"
        "Fine-tuned **DeBERTa-v3-large** ranking five candidate answers. "
        "Scored by MAP@3: credit 1, 1/2, 1/3 if the correct option is ranked 1st, 2nd or 3rd.\n\n"
        "*3-fold grouped CV MAP@3 **0.7567** — random baseline is 0.3667. "
        "Closed-book: it answers from its weights, with no retrieval and no abstention.*"
    )
    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(label="Question", lines=2,
                                  placeholder="Ask a science or philosophy question...")
            with gr.Row():
                a = gr.Textbox(label="A")
                b = gr.Textbox(label="B")
            with gr.Row():
                c = gr.Textbox(label="C")
                d = gr.Textbox(label="D")
            e = gr.Textbox(label="E")
            btn = gr.Button("Rank the options", variant="primary")
        with gr.Column(scale=2):
            verdict = gr.Markdown(label="Answer")
            table = gr.Dataframe(headers=["Option", "Text", "Probability", "Rank"],
                                 datatype=["str", "str", "number", "number"],
                                 label="All five, ranked", wrap=True)

    btn.click(solve, [question, a, b, c, d, e], [verdict, table])
    gr.Examples(EXAMPLES, inputs=[question, a, b, c, d, e])

    gr.Markdown(
        "---\n"
        "**Limitations.** Trained on 2,000 rows covering 252 unique questions, so coverage "
        "is narrow. It has no way to say *I don't know* — it will rank five options "
        "confidently even when it knows nothing about the topic. Reported scores reflect a "
        "dataset whose test split overlaps its training split heavily; the leakage-free "
        "estimate for this project is 0.6817."
    )

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