"""
app.py — Closed-World Grounding: live fabrication viewer
=========================================================================
An interactive Hugging Face Space that lets a judge SEE hallucination get
caught. It reads the benchmark's saved outputs (results/tier1_benchmark_detail.json)
and, for any variant, shows the ungrounded output next to the grounded one, with
every detected fabrication highlighted inline.

No API key, no live generation: the Space displays the exact model outputs that
produced the published benchmark numbers, so it loads instantly and never breaks.

Run locally:
    pip install gradio
    python app.py
On Hugging Face Spaces: set SDK = gradio and this file as the entrypoint.
"""

from __future__ import annotations

import json
import html
import re
from pathlib import Path

import gradio as gr

# -------------------------------------------------------------------------
# Data loading
# -------------------------------------------------------------------------

DETAIL_PATH = Path("results/tier1_benchmark_detail.json")

# Load the benchmark detail once at startup. Each entry:
#   {accession, arm, output, fabrication:{hgvs_c_fabrications, ungrounded_evidence_terms,
#    citation_misattributions, clinical_classification_leak}, calibration:{...}}
with open(DETAIL_PATH, encoding="utf-8") as f:
    _RAW = json.load(f)

# Reshape into {accession: {"grounded": entry, "ungrounded": entry}}
VARIANTS: dict[str, dict[str, dict]] = {}
for entry in _RAW:
    acc = entry["accession"]
    VARIANTS.setdefault(acc, {})[entry["arm"]] = entry

# Sorted list of accessions that have both arms (for the dropdown).
ACCESSIONS = sorted(a for a, arms in VARIANTS.items()
                    if "grounded" in arms and "ungrounded" in arms)

# Fabrication categories -> (label, css class, adjudication note)
FAB_FIELDS = {
    "hgvs_c_fabrications":       ("Invented HGVS c. notation", "fab-hgvs"),
    "ungrounded_evidence_terms": ("Wrong / invented evidence value", "fab-evidence"),
    "citation_misattributions":  ("Misattributed / invented citation", "fab-citation"),
    "clinical_classification_leak": ("Clinical P/LP over-assertion", "fab-clinical"),
}


# -------------------------------------------------------------------------
# Highlighting
# -------------------------------------------------------------------------

def _collect_spans(fab: dict) -> list[tuple[str, str, str]]:
    """Return (needle, css_class, label) for each fabrication string to mark."""
    spans = []
    for field, (label, css) in FAB_FIELDS.items():
        for item in fab.get(field, []):
            # Detector items can be phrases or "asserted X != source Y" notes.
            # Pull a searchable needle: the leading token(s) before " != " or ":".
            needle = re.split(r"\s+!=\s+|:\s", str(item))[0].strip()
            # For quoted claims like "...'functional stud'...", extract the quote.
            m = re.search(r"'([^']+)'", str(item))
            if m:
                needle = m.group(1)
            if needle:
                spans.append((needle, css, label))
    return spans


def highlight(text: str, fab: dict) -> str:
    """Return HTML with fabrication substrings wrapped in colored marks."""
    if not text:
        return "<em>(no output)</em>"
    escaped = html.escape(text)
    spans = _collect_spans(fab)
    # Longest needles first so nested matches don't get clobbered.
    for needle, css, label in sorted(spans, key=lambda s: -len(s[0])):
        esc_needle = html.escape(needle)
        if esc_needle and esc_needle in escaped:
            mark = (f'<mark class="{css}" title="{html.escape(label)}">'
                    f'{esc_needle}</mark>')
            escaped = escaped.replace(esc_needle, mark)
    return escaped.replace("\n", "<br>")


def _count_fab(fab: dict, adjudicated: bool, arm: str) -> int:
    """Total fabrications; if adjudicated, drop the grounded clinical false positives."""
    total = sum(len(fab.get(k, [])) for k in FAB_FIELDS)
    if adjudicated and arm == "grounded":
        total -= len(fab.get("clinical_classification_leak", []))
    return max(total, 0)


def _badge(n: int) -> str:
    if n == 0:
        return '<span class="badge badge-clean">✓ 0 fabrications</span>'
    unit = "fabrication" if n == 1 else "fabrications"
    return f'<span class="badge badge-dirty">⚠ {n} {unit}</span>'


# -------------------------------------------------------------------------
# View builder
# -------------------------------------------------------------------------

def render(accession: str, adjudicated: bool):
    pair = VARIANTS[accession]
    g, u = pair["grounded"], pair["ungrounded"]
    g_fab, u_fab = g["fabrication"], u["fabrication"]

    g_n = _count_fab(g_fab, adjudicated, "grounded")
    u_n = _count_fab(u_fab, adjudicated, "ungrounded")

    def panel(title, arm_class, badge_html, body_html, cal):
        cal_note = ""
        if cal:
            ok = cal.get("appropriate")
            lvl = cal.get("confidence_level", "?")
            chip = "cal-ok" if ok else "cal-bad"
            txt = "calibrated" if ok else "miscalibrated"
            cal_note = (f'<span class="cal {chip}">confidence: {txt} '
                        f'(evidence: {lvl})</span>')
        return (f'<div class="panel {arm_class}">'
                f'<div class="panel-head"><span class="panel-title">{title}</span>'
                f'{badge_html}</div>'
                f'<div class="cal-row">{cal_note}</div>'
                f'<div class="output">{body_html}</div></div>')

    left = panel("Ungrounded (baseline)", "arm-ungrounded", _badge(u_n),
                 highlight(u["output"], u_fab), u.get("calibration"))
    right = panel("Grounded (closed-world)", "arm-grounded", _badge(g_n),
                  highlight(g["output"], g_fab), g.get("calibration"))

    adj_note = ""
    if adjudicated and len(g_fab.get("clinical_classification_leak", [])) > 0:
        adj_note = ('<p class="adj-note">Adjudication on: '
                    f'{len(g_fab["clinical_classification_leak"])} grounded '
                    'clinical flag(s) here are detector false positives '
                    '(hedged/negated language) and are not counted.</p>')

    return f'<div class="grid">{left}{right}</div>{adj_note}'


def gene_of(accession: str) -> str:
    # gene isn't in detail.json; infer nothing — just show accession.
    return accession


# -------------------------------------------------------------------------
# CSS — the signature is the inline fabrication highlighting
# -------------------------------------------------------------------------

CSS = """
:root {
  --ink: #14181f; --paper: #fbfbf9; --line: #e3e2dc;
  --grounded: #1f7a5c; --ungrounded: #b64a2e;
  --hgvs: #fff1cc; --hgvs-b: #e0a500;
  --evidence: #ffe0e0; --evidence-b: #d1453b;
  --citation: #e6e0ff; --citation-b: #6a4bd1;
  --clinical: #ffd9c2; --clinical-b: #c85a24;
}
.gradio-container { max-width: 1180px !important; }
#hero { border-bottom: 2px solid var(--ink); padding: 4px 0 14px; margin-bottom: 10px; }
#hero h1 { font-family: ui-monospace, "SF Mono", Menlo, monospace;
  font-size: 1.5rem; letter-spacing: -0.02em; margin: 0 0 4px; color: var(--ink); }
#hero p { margin: 0; color: #5b6270; font-size: 0.95rem; }
.grid { display: grid; grid-template-columns: 1fr 1fr; gap: 14px; margin-top: 6px; }
@media (max-width: 820px){ .grid { grid-template-columns: 1fr; } }
.panel { border: 1px solid var(--line); border-radius: 10px; background: var(--paper);
  overflow: hidden; }
.panel-head { display: flex; align-items: center; justify-content: space-between;
  padding: 10px 14px; border-bottom: 1px solid var(--line); }
.panel-title { font-family: ui-monospace, monospace; font-weight: 600; font-size: 0.9rem; }
.arm-ungrounded .panel-head { border-top: 4px solid var(--ungrounded); }
.arm-grounded  .panel-head { border-top: 4px solid var(--grounded); }
.output { padding: 14px 16px; font-size: 0.9rem; line-height: 1.62; color: var(--ink);
  white-space: normal; max-height: 460px; overflow-y: auto; }
.badge { font-size: 0.72rem; font-weight: 700; padding: 3px 9px; border-radius: 999px;
  font-family: ui-monospace, monospace; }
.badge-clean { background: #e4f4ed; color: var(--grounded); }
.badge-dirty { background: #fbe6de; color: var(--ungrounded); }
.cal-row { padding: 0 14px; }
.cal { display: inline-block; margin-top: 8px; font-size: 0.72rem; padding: 2px 8px;
  border-radius: 6px; font-family: ui-monospace, monospace; }
.cal-ok { background: #e4f4ed; color: var(--grounded); }
.cal-bad { background: #fbe6de; color: var(--ungrounded); }
mark { padding: 0.5px 3px; border-radius: 3px; font-weight: 600; }
mark.fab-hgvs     { background: var(--hgvs);     box-shadow: inset 0 -2px var(--hgvs-b); }
mark.fab-evidence { background: var(--evidence); box-shadow: inset 0 -2px var(--evidence-b); }
mark.fab-citation { background: var(--citation); box-shadow: inset 0 -2px var(--citation-b); }
mark.fab-clinical { background: var(--clinical); box-shadow: inset 0 -2px var(--clinical-b); }
.adj-note { margin-top: 12px; font-size: 0.82rem; color: #5b6270; border-left: 3px solid
  var(--grounded); padding-left: 10px; }
.legend { display: flex; flex-wrap: wrap; gap: 12px; margin: 10px 0 2px; font-size: 0.75rem;
  font-family: ui-monospace, monospace; color: #5b6270; }
.legend span { display: inline-flex; align-items: center; gap: 5px; }
.swatch { width: 13px; height: 13px; border-radius: 3px; display: inline-block; }
"""

LEGEND = """
<div class="legend">
  <span><i class="swatch" style="background:#fff1cc;box-shadow:inset 0 -2px #e0a500"></i>Invented HGVS c.</span>
  <span><i class="swatch" style="background:#ffe0e0;box-shadow:inset 0 -2px #d1453b"></i>Wrong evidence value</span>
  <span><i class="swatch" style="background:#e6e0ff;box-shadow:inset 0 -2px #6a4bd1"></i>Misattributed citation</span>
  <span><i class="swatch" style="background:#ffd9c2;box-shadow:inset 0 -2px #c85a24"></i>Clinical P/LP over-assertion</span>
</div>
"""

HERO = """
<div id="hero">
  <h1>Closed-World Grounding — fabrication viewer</h1>
  <p>Same variant, same model, two prompts. The left arm runs unconstrained; the right
  is held to a closed world of verified facts. Highlights mark every fabrication the
  detectors caught — pick a variant and compare.</p>
</div>
"""


# -------------------------------------------------------------------------
# App
# -------------------------------------------------------------------------

with gr.Blocks(title="Closed-World Grounding") as demo:
    # CSS is attached here for broad Gradio-version compatibility (in Gradio 6
    # the Blocks(css=...) arg moved; injecting a <style> block works everywhere).
    gr.HTML(f"<style>{CSS}</style>")
    gr.HTML(HERO)
    with gr.Row():
        acc = gr.Dropdown(ACCESSIONS, value=ACCESSIONS[0] if ACCESSIONS else None,
                          label="Variant (VCV accession)", scale=3)
        adj = gr.Checkbox(value=True, label="Adjudicate detector false positives", scale=1)
    gr.HTML(LEGEND)
    out = gr.HTML()

    def _update(a, j):
        if not a:
            return "<em>No variant selected.</em>"
        return render(a, j)

    acc.change(_update, [acc, adj], out)
    adj.change(_update, [acc, adj], out)
    demo.load(_update, [acc, adj], out)

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