try:
    import spaces
except ImportError:
    pass

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForMultipleChoice
import torch.nn.functional as F

# Load model
MODEL_DIR = "./saved_model"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForMultipleChoice.from_pretrained(MODEL_DIR)
model.to(device)
model.eval()
print("✅ Model loaded and ready.")

OPTION_LABELS = ["A", "B", "C", "D", "E"]
MEDAL = {"A": "🥇", "B": "🥈", "C": "🥉", "D": "4️⃣", "E": "5️⃣"}

CUSTOM_CSS = """
/* ── Global ── */
body, .gradio-container {
    font-family: 'Inter', 'Segoe UI', sans-serif !important;
    background: #0f1117 !important;
}

/* ── Header card ── */
#header-card {
    background: linear-gradient(135deg, #1e3a5f 0%, #0d2137 50%, #1a1f35 100%);
    border-radius: 16px;
    padding: 32px 40px;
    margin-bottom: 24px;
    border: 1px solid rgba(99, 179, 237, 0.2);
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4);
}
#header-card h1 {
    font-size: 2.2rem !important;
    font-weight: 800 !important;
    background: linear-gradient(90deg, #63b3ed, #a78bfa, #f687b3);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0 0 8px 0 !important;
}
#header-card p {
    color: #a0aec0 !important;
    font-size: 0.95rem !important;
    margin: 0 !important;
}

/* ── Input/output panels ── */
.panel {
    background: #1a1f2e !important;
    border: 1px solid rgba(99, 179, 237, 0.15) !important;
    border-radius: 12px !important;
    padding: 20px !important;
}

/* ── Labels ── */
label span, .label-wrap span {
    color: #90cdf4 !important;
    font-weight: 600 !important;
    font-size: 0.85rem !important;
    text-transform: uppercase !important;
    letter-spacing: 0.05em !important;
}

/* ── Textareas & inputs ── */
textarea, input[type="text"] {
    background: #0f1117 !important;
    border: 1px solid rgba(99, 179, 237, 0.2) !important;
    border-radius: 8px !important;
    color: #e2e8f0 !important;
    font-size: 0.95rem !important;
    transition: border-color 0.2s ease !important;
}
textarea:focus, input[type="text"]:focus {
    border-color: #63b3ed !important;
    box-shadow: 0 0 0 2px rgba(99, 179, 237, 0.15) !important;
    outline: none !important;
}

/* ── Submit button ── */
#submit-btn {
    background: linear-gradient(135deg, #3182ce, #553c9a) !important;
    border: none !important;
    border-radius: 10px !important;
    color: white !important;
    font-size: 1rem !important;
    font-weight: 700 !important;
    padding: 14px 0 !important;
    width: 100% !important;
    cursor: pointer !important;
    transition: all 0.25s ease !important;
    box-shadow: 0 4px 15px rgba(49, 130, 206, 0.3) !important;
}
#submit-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(49, 130, 206, 0.5) !important;
    filter: brightness(1.1) !important;
}
#submit-btn:active {
    transform: translateY(0) !important;
}

/* ── Result boxes ── */
#top3-output textarea {
    font-size: 2rem !important;
    font-weight: 800 !important;
    text-align: center !important;
    color: #f6e05e !important;
    letter-spacing: 0.3em !important;
    background: #0f1117 !important;
    border: 1px solid rgba(246, 224, 94, 0.3) !important;
}
#breakdown-output textarea {
    font-family: 'JetBrains Mono', 'Fira Code', monospace !important;
    font-size: 0.9rem !important;
    color: #a0aec0 !important;
    line-height: 1.8 !important;
    background: #0f1117 !important;
}

/* ── Stats bar ── */
#stats-bar {
    background: #1a1f2e;
    border: 1px solid rgba(99, 179, 237, 0.1);
    border-radius: 10px;
    padding: 14px 24px;
    text-align: center;
    color: #718096;
    font-size: 0.82rem;
    margin-top: 16px;
    letter-spacing: 0.03em;
}
#stats-bar span { color: #63b3ed; font-weight: 600; }

/* ── Examples section ── */
.examples-holder table {
    background: #1a1f2e !important;
    border-radius: 8px !important;
    overflow: hidden !important;
}
.examples-holder td, .examples-holder th {
    color: #a0aec0 !important;
    border-color: rgba(99, 179, 237, 0.1) !important;
    font-size: 0.85rem !important;
}
"""


@spaces.GPU
def predict(prompt, opt_a, opt_b, opt_c, opt_d, opt_e):
    """Tokenize, run model forward pass, and return ranked answers."""
    if not prompt.strip():
        return "⚠️ Please enter a question.", ""

    options = [
        f"(A) {opt_a}",
        f"(B) {opt_b}",
        f"(C) {opt_c}",
        f"(D) {opt_d}",
        f"(E) {opt_e}",
    ]

    prompts = [prompt] * 5
    encoding = tokenizer(
        prompts,
        options,
        truncation=True,
        max_length=256,
        padding="max_length",
        return_tensors="pt",
    )

    input_ids = encoding["input_ids"].unsqueeze(0).to(device)
    attention_mask = encoding["attention_mask"].unsqueeze(0).to(device)
    token_type_ids = encoding.get("token_type_ids")
    if token_type_ids is not None:
        token_type_ids = token_type_ids.unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )

    logits = outputs.logits.squeeze(0)
    probs = F.softmax(logits, dim=-1).cpu().tolist()

    ranked = sorted(
        zip(OPTION_LABELS, probs), key=lambda x: x[1], reverse=True
    )

    top3_str = "  ".join([label for label, _ in ranked[:3]])
    medals = ["🥇", "🥈", "🥉", "4th", "5th"]
    breakdown_lines = []
    for i, (label, prob) in enumerate(ranked):
        bar_len = int(prob * 30)
        bar = "█" * bar_len + "░" * (30 - bar_len)
        breakdown_lines.append(
            f"{medals[i]}  {label}  [{bar}]  {prob*100:.2f}%"
        )
    breakdown_str = "\n".join(breakdown_lines)

    return top3_str, breakdown_str

# Build UI
with gr.Blocks(
    title="Smart MCQ Solver · DeBERTa-v3",
) as demo:
    
    gr.HTML(f"<style>{CUSTOM_CSS}</style>")

    gr.HTML("""
    <div id="header-card">
        <h1>🎓 Smart MCQ Solver</h1>
        <p>
            Powered by a fine-tuned <strong style="color:#90cdf4;">DeBERTa-v3-small</strong> model.
            Enter your question and five answer options — the model ranks the top 3 most likely answers.
        </p>
    </div>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=3, elem_classes="panel"):
            prompt_box = gr.Textbox(
                label="Question / Context Prompt",
                placeholder="e.g.  Pick the best possible answer: What is the capital of France?",
                lines=4,
                max_lines=10,
            )
            with gr.Row():
                opt_a = gr.Textbox(label="Option A", placeholder="Paris")
                opt_b = gr.Textbox(label="Option B", placeholder="London")
            with gr.Row():
                opt_c = gr.Textbox(label="Option C", placeholder="Berlin")
                opt_d = gr.Textbox(label="Option D", placeholder="Rome")
            with gr.Row():
                opt_e = gr.Textbox(label="Option E", placeholder="Madrid")

            submit_btn = gr.Button(
                "🔍  Rank My Answers",
                variant="primary",
                elem_id="submit-btn",
            )

        with gr.Column(scale=2, elem_classes="panel"):
            gr.Markdown("### 🏆 Top-3 Predictions")
            top3_out = gr.Textbox(
                label="Ranked Answer Letters",
                interactive=False,
                elem_id="top3-output",
            )
            gr.Markdown("### 📊 Confidence Breakdown")
            breakdown_out = gr.Textbox(
                label="All Options (ranked by confidence)",
                lines=7,
                interactive=False,
                elem_id="breakdown-output",
            )

    gr.Examples(
        examples=[
            [
                "Which of the following is correct? Who proposed the concept of 'maximal acceleration'? carefully.",
                "Max Planck",
                "Niels Bohr",
                "Eduardo R. Caianiello",
                "Hideki Yukawa",
                "Albert Einstein",
            ],
            [
                "Pick the best possible answer: What does CPU stand for?",
                "Central Processing Unit",
                "Computer Personal Unit",
                "Central Program Utility",
                "Control Processing Unit",
                "Core Processor Uniblock",
            ],
        ],
        inputs=[prompt_box, opt_a, opt_b, opt_c, opt_d, opt_e],
        outputs=[top3_out, breakdown_out],
        fn=predict,
        cache_examples=False,
        label="💡 Try Example Questions",
    )

    gr.HTML("""
    <div id="stats-bar">
        Model: <span>microsoft/deberta-v3-small</span> &nbsp;|&nbsp;
        Task: <span>Multiple-Choice QA</span> &nbsp;|&nbsp;
        Kaggle MAP@3: <span>0.75353</span> &nbsp;|&nbsp;
        Training epochs: <span>5</span> &nbsp;|&nbsp;
        Running on: <span>CPU / GPU (auto)</span>
    </div>
    """)

    submit_btn.click(
        fn=predict,
        inputs=[prompt_box, opt_a, opt_b, opt_c, opt_d, opt_e],
        outputs=[top3_out, breakdown_out],
    )

demo.launch()