"""Grounding demo — is a claim supported by the document it cites?

Runs BOTH grounding models side by side:
  * grounding-en              (open weights, English)
  * grounding-multilingual    (commercial; loaded from a PRIVATE repo via the HF_TOKEN Space secret,
                               so its outputs are shown but the weights are never downloadable here)

Each score is a calibrated support probability: p = softmax(logits / T)[entailment_index]. Temperature T
is per-model (monotonic — it fixes the confidence value, not the ranking).
"""
import os

import gradio as gr
import spaces
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

EI = 0  # entailment class index for the *-zeroshot-v2.0 heads (id2label {0: entailment, 1: not_entailment})

MODELS = [
    {"key": "en", "name": "grounding-en", "tag": "open · English",
     "repo": os.environ.get("EN_REPO", "nutrientdocs/grounding-en"), "T": 1.2866, "private": False},
    {"key": "multi", "name": "grounding-multilingual", "tag": "commercial · 15+ languages",
     "repo": os.environ.get("ML_REPO", "nutrientdocs/grounding-multilingual-private"),
     "subfolder": "weights", "T": 0.9432, "private": True},
]

TOKEN = os.environ.get("HF_TOKEN")  # Space secret; required for the private commercial model
_loaded = {}

# Model selector -> which model keys to score. A model card deep-links here with ?model=en|multi so
# clicking "Try it" from a model page opens the demo focused on that model.
SELECTION = {
    "grounding-en (English)": ["en"],
    "grounding-multilingual": ["multi"],
    "Both (compare)": ["en", "multi"],
}


def _from_query(request: gr.Request):
    m = (request.query_params.get("model") or "").lower()
    if m == "en":
        return "grounding-en (English)"
    if m in ("multi", "multilingual", "bge"):
        return "grounding-multilingual"
    return "Both (compare)"


def _ensure(spec):
    """Load tokenizer+model onto CPU and cache. Runs OUTSIDE the GPU window (no GPU needed to download),
    so a ZeroGPU allocation is spent only on the forward pass. Stores the Exception if a private model
    can't be reached (e.g. missing HF_TOKEN secret)."""
    if spec["key"] in _loaded:
        return _loaded[spec["key"]]
    kw = {"token": TOKEN} if TOKEN else {}  # both models are private; token is org-scoped
    if spec.get("subfolder"):
        kw["subfolder"] = spec["subfolder"]
    try:
        tok = AutoTokenizer.from_pretrained(spec["repo"], **kw)
        model = AutoModelForSequenceClassification.from_pretrained(spec["repo"], **kw).eval()
        _loaded[spec["key"]] = (tok, model)
    except Exception as e:
        _loaded[spec["key"]] = e
    return _loaded[spec["key"]]


@spaces.GPU(duration=120)
def _forward(premise, hypothesis, keys):
    """ZeroGPU-allocated: move each selected model to CUDA and score. Returns key -> calibrated support
    probability (None if the model failed to load)."""
    scores = {}
    for spec in MODELS:
        if spec["key"] not in keys:
            continue
        got = _loaded.get(spec["key"])
        if not isinstance(got, tuple):
            scores[spec["key"]] = None
            continue
        tok, model = got
        model = model.to("cuda")
        enc = tok(premise, hypothesis, truncation=True, max_length=1024, return_tensors="pt").to("cuda")
        with torch.no_grad():
            logits = model(**enc).logits
        scores[spec["key"]] = torch.softmax(logits / spec["T"], dim=-1)[0, EI].item()  # calibrated
    return scores


def grade(premise, hypothesis, selection):
    # Generator: emit an immediate status line so the click has visible feedback while the (potentially
    # slow) cold start runs — first request downloads the weights to CPU AND cold-starts a shared ZeroGPU
    # allocation, which together can take ~30-60s. Warm requests only wait on the GPU handoff.
    if not premise.strip() or not hypothesis.strip():
        yield "Enter a document premise and a claim to check."
        return
    keys = SELECTION.get(selection, ["en", "multi"])
    cold = [s for s in MODELS if s["key"] in keys and s["key"] not in _loaded]
    if cold:
        yield ("⏳ **Waking up…** the first check downloads the model and cold-starts a shared GPU — "
               "this can take **~30–60s**. Later checks are near-instant.")
    else:
        yield "⏳ Scoring on the GPU…"
    for spec in MODELS:            # download/load on CPU first, so the GPU window is inference-only
        if spec["key"] in keys:
            _ensure(spec)
    scores = _forward(premise, hypothesis, keys)
    lines = ["| Model | | Grounded support | Verdict |", "|---|---|---:|---|"]
    for spec in MODELS:
        if spec["key"] not in keys:
            continue
        p = scores.get(spec["key"])
        if p is None:
            lines.append(f"| **{spec['name']}** | {spec['tag']} | _unavailable_ | "
                         "_(commercial — set HF_TOKEN)_ |")
            continue
        verdict = "✅ grounded" if p >= 0.5 else "❌ not grounded"
        bar = "█" * round(p * 10) + "░" * (10 - round(p * 10))
        lines.append(f"| **{spec['name']}** | {spec['tag']} | `{bar}` **{p:.3f}** | {verdict} |")
    yield "\n".join(lines)


from examples import EXAMPLES  # noqa: E402  (real, labeled rows; single source of truth, verified)

with gr.Blocks(title="Grounding demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# Does the document actually support this claim?\n"
        "Paste a **document premise** (a table or passage) and a **claim**. The grounding model scores "
        "whether the document *supports* the claim. Scores are temperature-calibrated support "
        "probabilities; the verdict is the model's decision at 0.5.\n\n"
        "Pick an example below (its **expected answer** is shown) or type your own. The non-English "
        "examples are cases the English model gets wrong — switch to the **multilingual** model.\n\n"
        "→ [grounding-en (open)](https://huggingface.co/nutrientdocs/grounding-en) · "
        "[grounding-multilingual (commercial)](https://huggingface.co/nutrientdocs/grounding-multilingual) · "
        "[leaderboard](https://huggingface.co/spaces/nutrientdocs/grounding-leaderboard) · "
        "[benchmark](https://huggingface.co/datasets/nutrientdocs/grounding-benchmark)")
    model_sel = gr.Radio(list(SELECTION), value="Both (compare)", label="Model")
    with gr.Row():
        premise = gr.Textbox(label="Document premise", lines=8)
        hypothesis = gr.Textbox(label="Claim to check", lines=8)
    expected = gr.Textbox(label="Expected answer (ground truth for the selected example)",
                          interactive=False)
    btn = gr.Button("Check grounding", variant="primary")
    gr.Markdown("<sub>⏱️ The first check cold-starts a shared GPU (~30–60s); later checks are fast.</sub>")
    out = gr.Markdown()
    btn.click(grade, [premise, hypothesis, model_sel], out, show_progress="full")
    gr.Examples(EXAMPLES, [premise, hypothesis, expected])  # fills the boxes incl. expected answer
    gr.Markdown(
        '## About the author\n'
        '<a href="https://nutrient.io/">'
        '<img src="https://avatars2.githubusercontent.com/u/1527679?v=3&s=200" height="80" /></a>\n\n'
        "This demo is maintained and funded by [Nutrient](https://nutrient.io/) — "
        "The deterministic document infrastructure enterprises run their highest-stakes workflows on: "
        "replayable output, clear exceptions, and full audit trails on the messy, regulated documents where AI alone breaks.")
    demo.load(_from_query, None, model_sel)                 # ?model=en|multi presets the selector

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