"""Open-vocab document/figure classification demo — doc-img-classification (ZeroGPU).

Upload a document or figure image, type any set of candidate labels, and the model ranks them. Runs the
COMMERCIAL model from a PRIVATE repo via the HF_TOKEN Space secret: a pre-merged, base-scrubbed SigLIP2
checkpoint loaded with plain `from_pretrained` (no adapter, no base-model identity), scored on a ZeroGPU
allocation. Weights are used server-side only and are never downloadable.

Per candidate label: score = sigmoid((logit_scale·cos(image, label) + logit_bias) / T), the SigLIP
per-pair probability with the calibration temperature applied. Scores are independent per label (SigLIP
sigmoid), not a softmax across the set.
"""
import os

import gradio as gr
import spaces
import torch

REPO = os.environ.get("DOC_ADAPT_REPO", "nutrientdocs/doc-img-classification-private")
TOKEN = os.environ.get("HF_TOKEN")           # Space secret; required to reach the private weights
T = float(os.environ.get("CALIB_T", "0.745"))  # fitted calibration temperature (baselines/temperature.json)

DEFAULT_LABELS = ("invoice, resume, scientific publication, receipt, financial report, "
                  "bar chart, screenshot, signature, letter, form")

_state = {}                                  # lazy: {"model":, "processor":} or {"error": Exception}


def _pooled(feats):
    return feats.pooler_output if hasattr(feats, "pooler_output") else feats


def _load():
    """Load the merged model + processor onto CPU and cache (runs OUTSIDE the GPU window — no GPU needed
    to download). Stores the Exception if the private repo can't be reached (e.g. missing HF_TOKEN)."""
    if _state:
        return _state
    try:
        from transformers import AutoModel, AutoProcessor
        kw = {"token": TOKEN} if TOKEN else {}
        _state["model"] = AutoModel.from_pretrained(REPO, torch_dtype=torch.float32, **kw).eval()
        _state["processor"] = AutoProcessor.from_pretrained(REPO, **kw)
    except Exception as e:                    # noqa: BLE001 — surfaced in the UI (usually a missing secret)
        _state["error"] = e
    return _state


@spaces.GPU(duration=90)
def _score(image, labels):
    """ZeroGPU-allocated: embed the image + candidate labels on CUDA and return per-label sigmoid scores."""
    model, processor = _state["model"].to("cuda"), _state["processor"]
    enc_i = processor(images=[image], return_tensors="pt").to("cuda")
    enc_t = processor(text=labels, padding="max_length", max_length=64, truncation=True,
                      return_tensors="pt").to("cuda")
    with torch.no_grad():
        ie = torch.nn.functional.normalize(_pooled(model.get_image_features(**enc_i)), dim=-1)   # [1,D]
        te = torch.nn.functional.normalize(_pooled(model.get_text_features(**enc_t)), dim=-1)     # [C,D]
        logits = (ie @ te.t()) * model.logit_scale.exp() + model.logit_bias                       # [1,C]
        probs = torch.sigmoid(logits / T)[0]
    return probs.cpu().tolist()


def classify(image, labels_text):
    if image is None:
        return "Upload a document or figure image first."
    labels = [s.strip() for s in (labels_text or "").replace("\n", ",").split(",") if s.strip()]
    if not labels:
        return "Enter at least one candidate label (comma-separated)."
    st = _load()
    if "error" in st:
        return (f"⚠️ Model unavailable — could not load the commercial weights "
                f"(`{type(st['error']).__name__}`). This Space needs an **HF_TOKEN** secret with read "
                f"access to `{REPO}` (Settings → Variables and secrets).")
    scores = _score(image.convert("RGB"), labels)
    ranked = sorted(zip(labels, scores), key=lambda x: -x[1])
    lines = ["| Rank | Label | Score |", "|---:|---|---:|"]
    for i, (lab, p) in enumerate(ranked, 1):
        bar = "█" * round(p * 10) + "░" * (10 - round(p * 10))
        mark = " ✅" if i == 1 else ""
        lines.append(f"| {i} | **{lab}**{mark} | `{bar}` {p:.3f} |")
    return "\n".join(lines)


def _from_query(request: gr.Request):
    q = (request.query_params.get("labels") or "").strip()
    return q if q else DEFAULT_LABELS


try:
    from examples import EXAMPLES          # [[image_path, labels_text], ...] — hard docs we classify right
except Exception:                          # noqa: BLE001 — examples optional; demo works without them
    EXAMPLES = []


with gr.Blocks(title="Open-vocab document classification", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# Classify any document or figure against labels you choose\n"
        "Upload a **document page or figure image**, type a comma-separated list of **candidate labels**, "
        "and `doc-img-classification` ranks them — open-vocabulary, no fixed class list. Scores are calibrated "
        "per-label match probabilities.\n\n"
        "→ [model](https://huggingface.co/nutrientdocs/doc-img-classification) · "
        "[leaderboard](https://huggingface.co/spaces/nutrientdocs/doc-openvocab-leaderboard) · "
        "[benchmark](https://huggingface.co/datasets/nutrientdocs/doc-openvocab-benchmark)")
    with gr.Row():
        image = gr.Image(label="Document / figure image", type="pil", height=380)
        labels = gr.Textbox(label="Candidate labels (comma-separated)", value=DEFAULT_LABELS, lines=8)
    btn = gr.Button("Classify", variant="primary")
    gr.Markdown("<sub>⏱️ The first classify cold-starts a shared GPU (~20–40s); later ones are fast.</sub>")
    out = gr.Markdown()
    btn.click(classify, [image, labels], out, show_progress="full")
    if EXAMPLES:
        # cache_examples=False: caching would run the model at build time, when the ZeroGPU allocation +
        # the private weights (HF_TOKEN secret) aren't available. Clicking one fills the inputs; the user
        # then hits Classify.
        gr.Examples(EXAMPLES, inputs=[image, labels], cache_examples=False,
                    label="Try a hard document the model gets right (click, then Classify)")
    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, labels)                # ?labels=a,b,c presets the candidate list

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