src/app.py

"""
app.py — Streamlit frontend for the Smart MCQ Solver.

Run locally:
    streamlit run src/app.py
"""

import sys
import os
import numpy as np

# ── Ensure the project root is on the Python path so `src.*` imports work ────
_SRC_DIR = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.dirname(_SRC_DIR)
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

import streamlit as st

from src.config import (
    DEVICE, ID2LABEL, BAR_COLORS, EXAMPLES,
    MODEL_SOURCE,
)
from src.model import load_roberta_model, predict_mcq

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Smart MCQ Solver",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── Session State Init ────────────────────────────────────────────────────────
for k in ["q_prompt", "opt_A", "opt_B", "opt_C", "opt_D", "opt_E"]:
    if k not in st.session_state:
        st.session_state[k] = ""

def load_example():
    """Callback to update text areas when an example is selected."""
    choice = st.session_state.ex_choice_box
    if choice != "(none)":
        idx = int(choice.split()[-1]) - 1
        ex = EXAMPLES[idx]
        st.session_state.q_prompt = ex["prompt"]
        st.session_state.opt_A = ex["A"]
        st.session_state.opt_B = ex["B"]
        st.session_state.opt_C = ex["C"]
        st.session_state.opt_D = ex["D"]
        st.session_state.opt_E = ex["E"]
    else:
        st.session_state.q_prompt = ""
        st.session_state.opt_A = ""
        st.session_state.opt_B = ""
        st.session_state.opt_C = ""
        st.session_state.opt_D = ""
        st.session_state.opt_E = ""


# ── Custom CSS ────────────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

html, body, [class*="css"] {
    font-family: 'Inter', sans-serif;
}

.stApp {
    background: linear-gradient(135deg, #0f0c29, #302b63, #24243e);
    color: #e0e0e0;
}

/* ── Glass card ─────────────────────────────────────────────────────────── */
.glass-card {
    background: rgba(255, 255, 255, 0.05);
    border: 1px solid rgba(255, 255, 255, 0.1);
    border-radius: 16px;
    padding: 24px;
    backdrop-filter: blur(16px);
    margin-bottom: 16px;
    box-shadow: 0 4px 30px rgba(0, 0, 0, 0.1);
}

/* ── Hero ────────────────────────────────────────────────────────────────── */
.hero-title {
    font-size: 2.8rem;
    font-weight: 800;
    background: linear-gradient(90deg, #a78bfa, #60a5fa, #34d399);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    text-align: center;
    margin-bottom: 4px;
    letter-spacing: -0.02em;
}
.hero-sub {
    text-align: center;
    color: #94a3b8;
    font-size: 1.05rem;
    margin-bottom: 32px;
    font-weight: 400;
}

/* ── Prediction badge ───────────────────────────────────────────────────── */
.pred-badge {
    display: inline-block;
    font-size: 4rem;
    font-weight: 800;
    background: linear-gradient(135deg, #a78bfa, #60a5fa);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    line-height: 1.2;
}
.pred-label {
    font-size: 0.85rem;
    color: #64748b;
    text-transform: uppercase;
    letter-spacing: 0.12em;
    font-weight: 700;
}

/* ── Probability bars ───────────────────────────────────────────────────── */
.prob-row {
    display: flex;
    align-items: center;
    margin: 8px 0;
    gap: 12px;
}
.prob-label {
    font-size: 1.05rem;
    font-weight: 700;
    width: 28px;
    text-align: center;
    color: #e2e8f0;
}
.prob-track {
    flex: 1;
    height: 12px;
    background: rgba(255,255,255,0.06);
    border-radius: 12px;
    overflow: hidden;
}
.prob-fill {
    height: 100%;
    border-radius: 12px;
    transition: width 0.6s cubic-bezier(0.22, 1, 0.36, 1);
}
.prob-pct {
    font-size: 0.9rem;
    color: #94a3b8;
    font-weight: 500;
    width: 56px;
    text-align: right;
}

/* ── Chips ───────────────────────────────────────────────────────────────── */
.chip {
    display: inline-block;
    padding: 8px 18px;
    border-radius: 999px;
    font-weight: 600;
    font-size: 0.95rem;
    margin: 4px;
    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
}
.chip-1 { background: linear-gradient(135deg,#a78bfa,#7c3aed); color:#fff; }
.chip-2 { background: rgba(96,165,250,0.15); color:#60a5fa; border:1px solid rgba(96,165,250,0.4); }
.chip-3 { background: rgba(52,211,153,0.1); color:#34d399; border:1px solid rgba(52,211,153,0.3); }

/* ── Section headers ────────────────────────────────────────────────────── */
.section-header {
    font-size: 0.85rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.15em;
    color: #a78bfa;
    margin-bottom: 12px;
}

/* ── Text inputs ────────────────────────────────────────────────────────── */
.stTextArea textarea, .stTextInput input {
    background: rgba(255,255,255,0.04) !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    border-radius: 12px !important;
    color: #f1f5f9 !important;
    font-family: 'Inter', sans-serif !important;
    font-size: 0.95rem !important;
    line-height: 1.5 !important;
    padding: 12px 16px !important;
    transition: all 0.2s ease;
}
.stTextArea textarea:focus, .stTextInput input:focus {
    background: rgba(255,255,255,0.08) !important;
    border-color: #a78bfa !important;
    box-shadow: 0 0 0 3px rgba(167,139,250,0.2) !important;
}

/* ── Sidebar ────────────────────────────────────────────────────────────── */
section[data-testid="stSidebar"] {
    background: rgba(15, 12, 41, 0.95) !important;
    border-right: 1px solid rgba(255,255,255,0.05);
}

/* ── Button ─────────────────────────────────────────────────────────────── */
.stButton > button {
    background: linear-gradient(135deg, #8b5cf6, #3b82f6) !important;
    color: white !important;
    border: none !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
    font-size: 1.1rem !important;
    padding: 0.8rem 2rem !important;
    width: 100% !important;
    transition: all 0.2s ease !important;
    letter-spacing: 0.02em;
    box-shadow: 0 4px 14px rgba(59, 130, 246, 0.3) !important;
}
.stButton > button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(59, 130, 246, 0.4) !important;
}
.stButton > button:active {
    transform: translateY(0px) !important;
}

/* ── Selectbox ──────────────────────────────────────────────────────────── */
.stSelectbox > div > div {
    background: rgba(255,255,255,0.05) !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    border-radius: 10px !important;
    color: #e2e8f0 !important;
}

/* ── Metric card ────────────────────────────────────────────────────────── */
.metric-card {
    background: rgba(255,255,255,0.03);
    border: 1px solid rgba(255,255,255,0.08);
    border-radius: 12px;
    padding: 16px;
    text-align: center;
    transition: transform 0.2s ease;
}
.metric-card:hover {
    transform: translateY(-2px);
    background: rgba(255,255,255,0.05);
}
.metric-value {
    font-size: 1.4rem;
    font-weight: 700;
    color: #e2e8f0;
}
.metric-desc {
    font-size: 0.75rem;
    color: #94a3b8;
    margin-top: 6px;
    text-transform: uppercase;
    letter-spacing: 0.05em;
    font-weight: 600;
}

/* ── Status pill ────────────────────────────────────────────────────────── */
.status-pill {
    display: inline-block;
    padding: 6px 14px;
    border-radius: 999px;
    font-size: 0.75rem;
    font-weight: 700;
    letter-spacing: 0.05em;
}
.status-ok   { background: rgba(52,211,153,0.1); color: #34d399; border: 1px solid rgba(52,211,153,0.3); }
.status-warn { background: rgba(251,191,36,0.1);  color: #fbbf24; border: 1px solid rgba(251,191,36,0.3); }

</style>
""", unsafe_allow_html=True)


# ── HTML render helpers ───────────────────────────────────────────────────────

def render_prob_bars(probs: np.ndarray, highlight_idx: int) -> str:
    html = ""
    for i, p in enumerate(probs):
        width = max(int(p * 100), 1)
        color = BAR_COLORS[i]
        bold  = "font-weight:800;" if i == highlight_idx else "font-weight:600;"
        label_color = color if i == highlight_idx else "#94a3b8"
        html += f"""
        <div class="prob-row">
          <span class="prob-label" style="color:{label_color};{bold}">{ID2LABEL[i]}</span>
          <div class="prob-track">
            <div class="prob-fill" style="width:{width}%;background:{color};"></div>
          </div>
          <span class="prob-pct">{p*100:.1f}%</span>
        </div>"""
    return html


def render_top3_chips(probs: np.ndarray) -> str:
    top3 = np.argsort(-probs)[:3]
    chip_classes = ["chip chip-1", "chip chip-2", "chip chip-3"]
    rank_labels  = ["🥇", "🥈", "🥉"]
    html = ""
    for rank, idx in enumerate(top3):
        html += f'<span class="{chip_classes[rank]}">{rank_labels[rank]} {ID2LABEL[idx]} &nbsp;{probs[idx]*100:.1f}%</span>'
    return html


# ─────────────────────────────────────────────────────────────────────────────
# SIDEBAR
# ─────────────────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown("### 📋 Quick Examples")
    st.selectbox(
        "Load an example question",
        options=["(none)"] + [f"Example {i+1}" for i in range(len(EXAMPLES))],
        key="ex_choice_box",
        on_change=load_example,
        label_visibility="collapsed"
    )

    st.divider()
    st.markdown("### ℹ️ About the Model")
    source_label = "HF Hub" if MODEL_SOURCE == "hub" else "Local"
    hw_label = "🟢 GPU" if DEVICE == "cuda" else "🟡 CPU"
    
    st.markdown(f"""
    <div style="display:flex; flex-direction:column; gap:12px; margin-bottom:20px;">
      <div class="metric-card">
        <div class="metric-value">RoBERTa Base</div>
        <div class="metric-desc">Architecture</div>
      </div>
      <div style="display:flex; gap:12px;">
        <div class="metric-card" style="flex:1;">
          <div class="metric-value">r=8</div>
          <div class="metric-desc">LoRA Rank</div>
        </div>
        <div class="metric-card" style="flex:1;">
          <div class="metric-value">3.4MB</div>
          <div class="metric-desc">Adapter</div>
        </div>
      </div>
    </div>
    <div style="text-align:center;">
        <span class="status-pill {'status-ok' if DEVICE=='cuda' else 'status-warn'}">
        {hw_label} &nbsp; {DEVICE.upper()}
        </span>
    </div>
    """, unsafe_allow_html=True)


# ─────────────────────────────────────────────────────────────────────────────
# MAIN PAGE
# ─────────────────────────────────────────────────────────────────────────────
st.markdown('<div class="hero-title">🧠 Smart MCQ Solver</div>', unsafe_allow_html=True)
st.markdown(
    '<div class="hero-sub">Powered by LoRA fine-tuned RoBERTa · '
    'IITM DL+GenAI T2-2026</div>',
    unsafe_allow_html=True,
)

# ── Load model (cached across reruns) ─────────────────────────────────────────
@st.cache_resource(show_spinner=False)
def _cached_load():
    return load_roberta_model()

with st.spinner("⏳ Loading model weights into memory..."):
    try:
        rob_model, rob_tok = _cached_load()
    except Exception as err:
        st.error(f"❌ Failed to load model: {err}")
        st.stop()


# ─── Two-column layout ────────────────────────────────────────────────────────
left, right = st.columns([1.1, 0.9], gap="large")

# ── INPUT COLUMN ──────────────────────────────────────────────────────────────
with left:
    st.markdown('<div class="glass-card">', unsafe_allow_html=True)
    st.markdown('<div class="section-header">📝 The Question</div>', unsafe_allow_html=True)

    st.text_area(
        "Question / Prompt",
        height=110,
        placeholder="Enter your multiple-choice question here...",
        label_visibility="collapsed",
        key="q_prompt",
    )
    st.markdown("</div>", unsafe_allow_html=True)

    # ── Options ───────────────────────────────────────────────────────────────
    st.markdown('<div class="glass-card">', unsafe_allow_html=True)
    st.markdown('<div class="section-header">🔤 Options (A–E)</div>', unsafe_allow_html=True)

    for lbl in ["A", "B", "C", "D", "E"]:
        st.text_area(
            f"Option {lbl}",
            height=68,
            placeholder=f"Type option {lbl}...",
            key=f"opt_{lbl}",
        )

    st.markdown("</div>", unsafe_allow_html=True)

    predict_clicked = st.button("✨ Predict Answer", use_container_width=True)


# ── RESULT COLUMN ─────────────────────────────────────────────────────────────
with right:
    if predict_clicked:
        # Validation
        prompt = st.session_state.q_prompt.strip()
        opts = {
            "A": st.session_state.opt_A.strip(),
            "B": st.session_state.opt_B.strip(),
            "C": st.session_state.opt_C.strip(),
            "D": st.session_state.opt_D.strip(),
            "E": st.session_state.opt_E.strip(),
        }
        
        missing = []
        if not prompt: missing.append("Question")
        for k, v in opts.items():
            if not v: missing.append(f"Option {k}")

        if missing:
            st.error(f"⚠️ Please fill in: **{', '.join(missing)}**")
            st.stop()

        with st.spinner("🧠 Thinking..."):
            result = predict_mcq(prompt, opts, rob_model, rob_tok)

        probs       = result["probs"]
        top1_letter = result["top1"]
        top1_conf   = result["confidence"] * 100

        # ── Top-1 prediction card ─────────────────────────────────────────────
        st.markdown(f"""
        <div class="glass-card" style="text-align:center; padding: 40px 20px;">
          <div class="pred-label">Top Prediction</div>
          <div class="pred-badge">{top1_letter}</div>
          <div style="color:#94a3b8; margin-top:8px; font-size:1.05rem;">
            Confidence: <b style="color:#a78bfa">{top1_conf:.1f}%</b>
          </div>
        </div>
        """, unsafe_allow_html=True)

        # ── Top-3 chips ───────────────────────────────────────────────────────
        st.markdown('<div class="glass-card">', unsafe_allow_html=True)
        st.markdown('<div class="section-header">🏆 Top-3 Ranking</div>', unsafe_allow_html=True)
        st.markdown(f"<div style='margin-top:12px;'>{render_top3_chips(probs)}</div>", unsafe_allow_html=True)
        st.markdown("</div>", unsafe_allow_html=True)

        # ── Probability bars ──────────────────────────────────────────────────
        top1_idx = int(np.argmax(probs))
        st.markdown('<div class="glass-card">', unsafe_allow_html=True)
        st.markdown('<div class="section-header">📊 Detailed Probabilities</div>', unsafe_allow_html=True)
        st.markdown(f"<div style='margin-top:16px;'>{render_prob_bars(probs, top1_idx)}</div>", unsafe_allow_html=True)
        st.markdown("</div>", unsafe_allow_html=True)

    else:
        st.markdown("""
        <div class="glass-card" style="text-align:center; padding:80px 28px; display:flex; flex-direction:column; justify-content:center; align-items:center; height:100%;">
          <div style="font-size:3.5rem; margin-bottom:20px; filter: drop-shadow(0 0 10px rgba(167,139,250,0.5));">🎯</div>
          <div style="color:#e2e8f0; font-size:1.1rem; font-weight:500; margin-bottom:8px;">Ready for Inference</div>
          <div style="color:#94a3b8; font-size:0.95rem; line-height:1.6; max-width:250px;">
            Enter your question and options on the left, then click <b style="color:#a78bfa;">Predict Answer</b>.
          </div>
          <div style="margin-top:24px; padding-top:24px; border-top:1px solid rgba(255,255,255,0.1); color:#64748b; font-size:0.85rem;">
            Need a test? Load an example from the sidebar.
          </div>
        </div>
        """, unsafe_allow_html=True)


# ── Footer ────────────────────────────────────────────────────────────────────
st.markdown("""
<div style="text-align:center; color:#475569; font-size:0.8rem; padding: 24px 0; margin-top:20px; border-top:1px solid rgba(255,255,255,0.05);">
  Developed for <b>IITM DL+GenAI T2-2026</b>
</div>
""", unsafe_allow_html=True)