"""
Temporal Cell State Prediction — interactive demo (HF Space).

Frozen V-JEPA-2 ViT-L + a small attentive-pool head predicts per-cell cell-cycle
state (interphase / pre-mitosis / mitosis) directly from short single-cell clips,
replacing the classify stage of the segment->track->classify pipeline.

The Space serves PRE-COMPUTED results (no heavy model load on the free CPU tier):
a classification gallery (successes + failure modes), state-overlay + tracking videos,
the per-state counting probe, and a zero-shot VLM point-reasoning baseline.

Thesis of the study: in this small-data regime, *data scaling*, not model scaling,
is the binding constraint.
"""

from __future__ import annotations

import json
from pathlib import Path

import gradio as gr

ASSETS = Path(__file__).parent / "assets"


def _load_json(name: str, default):
    p = ASSETS / name
    return json.loads(p.read_text()) if p.exists() else default


METRICS = _load_json("metrics.json", {})
GALLERY = _load_json("gallery/gallery_manifest.json", [])


# ── Header / thesis ───────────────────────────────────────────────────────────
HEADER = """
# 🔬 Temporal State Prediction
**Frozen V-JEPA-2 ViT-L + attentive-pool head → per-cell cell-cycle state**, straight from a tracked clip —
no segment→track→**classify** hand-off.

> **Finding:** a frozen video foundation model is competitive with a purpose-built
> morphology–temporal baseline, and the binding constraint is **data scale, not model capacity.**
"""

LABELS3 = ["interphase", "pre-mitosis", "mitosis"]


def _selector_items():
    """Neutral clip thumbnails for the bottom selector (no prediction shown)."""
    items = []
    for g in GALLERY:
        img = ASSETS / "gallery" / g.get("raw_gif", g.get("gif", ""))
        if img.exists():
            items.append((str(img), f"cell {g.get('cell_id','?')} · t={g.get('t','?')}"))
    return items


def _valid_indices():
    return [i for i, g in enumerate(GALLERY)
            if (ASSETS / "gallery" / g.get("raw_gif", g.get("gif", ""))).exists()]


def _clip_view(i: int):
    """Return (image_path, result_markdown) for gallery item i."""
    g = GALLERY[i]
    img = str(ASSETS / "gallery" / g.get("raw_gif", g.get("gif", "")))
    verdict = "✅ **correct**" if g.get("correct") else f"❌ **incorrect** — true state is **{g['true']}**"
    probs = g.get("probs", [])
    bar = " · ".join(f"{LABELS3[k]} {probs[k]:.2f}" for k in range(len(probs))) if probs else ""
    note = ""
    if not g.get("correct"):
        if g["true"] == "pre-mitosis" and g["pred"] == "interphase":
            note = "\n\n*Why it fails: 'pre-mitosis' is a soft, lineage-defined window — this nucleus is morphologically identical to interphase (not yet rounded/condensed).*"
        elif g["true"] == "mitosis":
            note = "\n\n*A missed division — the rare, high-value event the metric weights most.*"
        elif g["pred"] == "pre-mitosis":
            note = "\n\n*Over-called: interphase nucleus flagged as entering division.*"
    md = (f"### Prediction: **{g['pred']}**  (confidence {g['conf']:.2f})\n\n"
          f"{verdict}\n\n"
          f"Per-class probability — {bar}{note}")
    return img, md


def _metrics_md() -> str:
    m = METRICS
    if not m:
        return "_metrics.json not found_"
    cm = m.get("confusion")
    lines = [
        "### Held-out HeLa (sequence 02, n=%s)" % m.get("n", "?"),
        "",
        "| model | macro-F1 | mitosis F1 | mitosis event P/R (±3fr) |",
        "|---|---|---|---|",
        f"| U-Net+BiLSTM baseline (3.8M) | {m.get('baseline_macro_f1','?')} | {m.get('baseline_mitosis_f1','?')} | {m.get('baseline_mit_pr','?')} |",
        f"| **frozen V-JEPA-2 head-only** | **{m.get('vjepa_macro_f1','?')}** | {m.get('vjepa_mitosis_f1','?')} | {m.get('vjepa_mit_pr','?')} |",
        "",
        f"*Data scaling (GOWT1→HeLa) lifts the same baseline +0.186 macro-F1; ~80× model scaling adds only +0.046.*",
        f"*Seed band: {m.get('seed_band','0.635 ± 0.098')} — single-seed gaps <0.08 are not significant.*",
    ]
    if cm:
        lines += [
            "",
            "**Confusion matrix** (rows = true, cols = pred):",
            "",
            "| true ⧵ pred | interphase | pre-mitosis | mitosis | recall |",
            "|---|---|---|---|---|",
        ]
        names = ["interphase", "pre-mitosis", "mitosis"]
        for i, nm in enumerate(names):
            row = cm[i]; rec = row[i] / max(sum(row), 1)
            lines.append(f"| **{nm}** | {row[0]} | {row[1]} | {row[2]} | {rec:.2f} |")
        lines.append("")
        lines.append("*Dominant error: pre-mitosis→interphase — a soft, lineage-defined 8-frame window with no sharp morphological boundary.*")
    return "\n".join(lines)


def build():
    with gr.Blocks(title="Temporal State Prediction", theme=gr.themes.Soft()) as demo:
        gr.Markdown(HEADER)
        with gr.Tabs():
            with gr.Tab("① Classify a clip"):
                gr.Markdown("**Select a single-cell clip below** — the model classifies that one clip into its cell-cycle state.")
                valid = _valid_indices()
                first = valid[0] if valid else 0
                img0, md0 = _clip_view(first) if valid else (None, "_no clips found_")
                with gr.Row():
                    sel_clip = gr.Image(value=img0, label="selected clip", height=300)
                    sel_md = gr.Markdown(md0)
                selector = gr.Gallery(value=_selector_items(), columns=5, height=200,
                                      object_fit="contain", label="▼ pick a clip",
                                      allow_preview=False)

                def _on_select(evt: gr.SelectData):
                    return _clip_view(evt.index)
                selector.select(_on_select, inputs=None, outputs=[sel_clip, sel_md])

                with gr.Accordion("Held-out test-set metrics (all 5,312 clips)", open=False):
                    gr.Markdown(_metrics_md())

            with gr.Tab("② Videos"):
                gr.Markdown("Whole field-of-view over the held-out sequence.")
                with gr.Row():
                    so = ASSETS / "videos" / "state_overlay.mp4"
                    ti = ASSETS / "videos" / "trackid.mp4"
                    if so.exists():
                        gr.Video(str(so), label="Predicted cell-cycle state (blue=interphase, amber=pre-mitosis, red=mitosis)", autoplay=True)
                    if ti.exists():
                        gr.Video(str(ti), label="Trackastra-style tracking (colour = track ID)", autoplay=True)

            with gr.Tab("③ VLM point-reasoning (zero-shot)"):
                gr.Markdown("A frontier VLM prompted to **point at each nucleus while reasoning**, then sum per state — "
                            "the 'visual primitives' recipe. Zero-shot baseline (no fine-tune); OOD-limited, shown for interpretability.")
                vlm = ASSETS / "vlm" / "vlm_overlay.png"
                if vlm.exists():
                    gr.Image(str(vlm), label="VLM predicted points (○) vs ground-truth centroids (×)")
                gr.Markdown((ASSETS / "vlm" / "vlm_trace.md").read_text() if (ASSETS / "vlm" / "vlm_trace.md").exists() else "")

        gr.Markdown("---\nModels: `DnaRnaProteins/vjepa2-cell-cycle-vit-l`, `DnaRnaProteins/unet-bilstm-cell-cycle-baseline` · "
                    "Data: MICCAI Cell Tracking Challenge (Fluo-N2DL-HeLa). Labels derived from lineage trees (no manual annotation).")
    return demo


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