"""Beyond Tokens · Decoding Playground — contrastive, guided and parallel decoding, step by step."""

import sys

print(f"[startup] python {sys.version.split()[0]}", flush=True)  # progress markers: see where a stuck start hangs

import spaces  # noqa: F401,E402  (ZeroGPU: must be imported before torch)

print("[startup] spaces imported", flush=True)

import json  # noqa: E402
import logging  # noqa: E402
import os  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402

print("[startup] gradio imported", flush=True)

import models  # noqa: E402

print("[startup] torch and models module imported", flush=True)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logging.getLogger("httpx").setLevel(logging.WARNING)
if os.environ.get("DD_SKIP_INIT") != "1":  # scripts/smoke_test.py imports this module only for its renderers
    models.init()

from content import examples as EX  # noqa: E402
from content import overview  # noqa: E402
from gpu_entry import gpu_run  # noqa: E402
from params import DEFAULT_REVERSE  # noqa: E402
from viz import inspector, parallel_viz, plots, tokens  # noqa: E402
from viz.css import CSS  # noqa: E402

LATEX = [{"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}]
GALLERY = Path(__file__).resolve().parent / "data" / "gallery"
FAMILY_CHOICES = [(fam.dropdown_label(), key) for key, fam in models.REGISTRY.items()]
DEFAULT_FAMILY = "qwen3" if "qwen3" in models.REGISTRY else (FAMILY_CHOICES[0][1] if FAMILY_CHOICES else None)

THEME = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="violet",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
)


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def call_gpu(paradigm: str, family: str, P: dict) -> dict:
    if family not in models.REGISTRY:
        raise gr.Error("No model family is available right now.")
    try:
        return gpu_run(paradigm, family, P)
    except ValueError as exc:
        raise gr.Error(str(exc)) from None


def compact(trace: dict | None) -> dict | None:
    if not trace:
        return None
    return {k: trace.get(k) for k in ("paradigm", "method", "family", "models", "params", "summary", "env")}


def trace_file(trace: dict) -> str:
    fd, path = tempfile.mkstemp(prefix=f"trace-{trace['paradigm']}-{trace['method']}-", suffix=".json")
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(trace, f, ensure_ascii=False)
    return path


def load_gallery(paradigm: str) -> dict | None:
    path = GALLERY / f"{paradigm}.json"
    try:
        return json.loads(path.read_text(encoding="utf-8")) if path.exists() else None
    except (OSError, json.JSONDecodeError):
        return None


def html(value="", **kwargs) -> gr.HTML:
    """gr.HTML without Gradio's default content styles, which in SSR mode override the panel CSS."""
    return gr.HTML(value, apply_default_css=False, **kwargs)


def strip_box(label: str, colors: dict, legend: bool = False, visible: bool = True) -> gr.HighlightedText:
    return gr.HighlightedText(label=label, color_map=colors, combine_adjacent=False, show_legend=legend,
                              show_inline_category=False, visible=visible)


def select_to_step(evt: gr.SelectData):
    idx = tokens.select_index(evt)
    return gr.update(value=idx) if idx is not None else gr.update()


def decoding_controls(max_new: int, cap: int, rep: float = 1.0):
    with gr.Accordion("Settings", open=False):
        with gr.Row():
            mode = gr.Radio([("greedy", "greedy"), ("sampling", "sampling")], value="greedy", label="Decoding")
            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.05, label="Temperature")
        with gr.Row():
            max_new_c = gr.Slider(8, cap, value=max_new, step=8, label="Max tokens")
            seed = gr.Number(value=0, precision=0, minimum=0, label="Seed")
        rep_c = gr.Slider(1.0, 1.5, value=rep, step=0.05, label="Repetition penalty")
        raw = gr.Checkbox(value=False, label="Raw prompt (no chat template)")
    return mode, temperature, max_new_c, seed, rep_c, raw


def trace_box():
    """Collapsed JSON summary plus a download of the full trace."""
    with gr.Accordion("Trace", open=False):
        summary = gr.JSON(show_label=False)
        download = gr.DownloadButton("Download", visible=False, size="sm")
    return summary, download


def slider_update(n: int, value: int) -> dict:
    return gr.update(minimum=0, maximum=max(n - 1, 1), value=max(0, min(value, max(n - 1, 0))))


# ---------------------------------------------------------------------------
# ① Contrastive
# ---------------------------------------------------------------------------


def run_contrastive(family, method, prompt, context, reverse, alpha, beta, tau, bucket, norm,
                    mode, temperature, max_new, seed, rep, raw):
    P = dict(method=method, prompt=prompt, context=context, reverse_system=reverse, alpha=alpha, beta=beta, tau=tau,
             bucket=bucket, apply_norm=norm, mode=mode, temperature=temperature, max_new_tokens=max_new, seed=seed,
             rep=rep, raw=raw)
    return call_gpu("contrastive", family, P)


def render_contrastive(trace):
    if not trace:
        return ("", gr.update(value=[]), gr.update(value=[]), gr.update(value=[], visible=False), slider_update(1, 0),
                inspector.empty(), gr.update(value="", visible=False), None, gr.update(visible=False))
    runs = trace["runs"]
    div = tokens.divergence(runs["baseline"], runs["method"])
    steps = runs["method"]["steps"]
    first = next((s["i"] for s in steps if s["changed"]), 0)
    third = runs.get("third")
    base_label, meth_label, third_label = tokens.labels(trace)
    return (
        parallel_viz.cards(trace),
        gr.update(value=tokens.strip_plain(runs["baseline"], div), label=base_label),
        gr.update(value=tokens.strip_steps(runs["method"], tokens.CHANGED_LABEL["contrastive"]), label=meth_label),
        gr.update(value=tokens.strip_plain(third) if third else [], label=third_label or "", visible=bool(third)),
        slider_update(len(steps), first),
        inspector.contrastive_step(trace, first),
        gr.update(value=inspector.dola_overview(trace, first) if trace["method"] == "dola" else "",
                  visible=trace["method"] == "dola"),
        compact(trace),
        gr.update(value=trace_file(trace), visible=True),
    )


def step_contrastive(trace, i):
    if not trace:
        return inspector.empty(), gr.update()
    i = int(i or 0)
    dola = gr.update(value=inspector.dola_overview(trace, i)) if trace["method"] == "dola" else gr.update()
    return inspector.contrastive_step(trace, i), dola


def replay_steps(trace, label):
    if not trace:
        return
    value = tokens.strip_steps(trace["runs"]["method"], label)
    for i in range(1, len(value) + 1):
        yield gr.update(value=value[:i])
        time.sleep(0.035)


def replay_contrastive(trace):
    yield from replay_steps(trace, tokens.CHANGED_LABEL["contrastive"])


def contrastive_method_changed(method, alpha, beta):
    return (
        gr.update(visible=method == "cad"),
        gr.update(visible=method == "rose"),
        gr.update(visible=method != "dola"),
        gr.update(visible=method == "cd"),
        gr.update(visible=method == "dola"),
        gr.update(visible=method == "dola"),
        gr.update(value=1.2 if method == "dola" else 1.0),
        overview.method_card(method, {"alpha": alpha, "beta": beta}),
    )


def build_contrastive(family):
    ex = EX.CONTRASTIVE[0]
    with gr.Row(equal_height=False):
        with gr.Column(scale=4, min_width=330):
            method = gr.Radio([("CD", "cd"), ("CAD", "cad"), ("ROSE", "rose"), ("DoLa", "dola")],
                              value=ex["method"], label="Method")
            prompt = gr.Textbox(ex["prompt"], lines=3, label="Prompt")
            context = gr.Textbox(ex["context"], lines=3, label="Context (z⁺ only)")
            reverse = gr.Textbox(DEFAULT_REVERSE, lines=2, label="Reverse system prompt (z⁻)", visible=False)
            alpha = gr.Slider(0, 3, value=ex["alpha"], step=0.05, label="α · contrast strength")
            beta = gr.Slider(0, 0.5, value=ex["beta"], step=0.01, label="β · plausibility cutoff")
            tau = gr.Slider(0.25, 2.0, value=1.0, step=0.05, label="τ · amateur temperature", visible=False)
            bucket = gr.Radio([("upper", "high"), ("lower", "low")], value="high", label="Candidate layers", visible=False)
            norm = gr.Checkbox(True, label="Final norm on early exits", visible=False)
            mode, temperature, max_new, seed, rep, raw = decoding_controls(48, 128)
            run = gr.Button("Run", variant="primary")
        with gr.Column(scale=7, min_width=420):
            note = gr.Markdown(overview.method_card(ex["method"], ex), latex_delimiters=LATEX)
            cards = html()
            base_ht = strip_box("Baseline", tokens.COLORS_BASELINE)
            meth_ht = strip_box("Method", tokens.COLORS_CONTRASTIVE, legend=True)
            third_ht = strip_box("z⁻ alone", tokens.COLORS_BASELINE)
            with gr.Row(equal_height=True):
                step = gr.Slider(0, 1, value=0, step=1, label="Step", scale=6)
                replay = gr.Button("▶ Replay", scale=1, min_width=90)
            insp = html(inspector.empty())
            dola = html(visible=False)
            summary, download = trace_box()
    inputs = [method, prompt, context, reverse, alpha, beta]
    gr.Examples(
        examples=[[e["method"], e["prompt"], e.get("context", ""), e.get("reverse_system", DEFAULT_REVERSE),
                   e["alpha"], e["beta"]] for e in EX.CONTRASTIVE],
        inputs=inputs, example_labels=[e["label"] for e in EX.CONTRASTIVE], label="Examples",
    )
    with gr.Accordion("Papers", open=False):
        html(overview.papers_html("contrastive"))

    state = gr.State(None)
    outputs = [cards, base_ht, meth_ht, third_ht, step, insp, dola, summary, download]
    method.change(contrastive_method_changed, [method, alpha, beta],
                  [context, reverse, alpha, tau, bucket, norm, rep, note], queue=False)
    for s in (alpha, beta):
        s.change(lambda m, a, b: overview.method_card(m, {"alpha": a, "beta": b}), [method, alpha, beta], note, queue=False)
    run.click(run_contrastive,
              [family, method, prompt, context, reverse, alpha, beta, tau, bucket, norm, mode, temperature, max_new, seed, rep, raw],
              state, concurrency_limit=4, concurrency_id="gpu", api_visibility="private"
              ).success(render_contrastive, state, outputs, queue=False)
    step.change(step_contrastive, [state, step], [insp, dola], queue=False, show_progress="hidden")
    meth_ht.select(select_to_step, None, step, queue=False)
    replay.click(replay_contrastive, state, meth_ht, concurrency_limit=None, api_visibility="private")
    return state, render_contrastive, outputs


# ---------------------------------------------------------------------------
# ② Guided
# ---------------------------------------------------------------------------


def run_guided(family, method, prompt, attribute, target, topic, lam, k, lookahead, pattern, top_k,
               mode, temperature, max_new, seed, rep, raw):
    P = dict(method=method, prompt=prompt, attribute=attribute, target=target, topic=topic, lam=lam, k=k,
             lookahead=lookahead, pattern=pattern, top_k=top_k, mode=mode, temperature=temperature,
             max_new_tokens=max_new, seed=seed, rep=rep, raw=raw)
    return call_gpu("guided", family, P)


def render_guided(trace):
    if not trace:
        return ("", gr.update(value=[]), gr.update(value=[]), slider_update(1, 0), inspector.empty(),
                gr.update(value=None), None, gr.update(visible=False))
    runs = trace["runs"]
    kind = trace["method"]
    steps = runs["method"]["steps"]
    first = next((s["i"] for s in steps if s["changed"]), 0)
    div = tokens.divergence(runs["baseline"], runs["method"])
    insp = inspector.classifier_step(trace, first) if kind == "classifier" else inspector.regex_step(trace, first)
    base_label, meth_label, _ = tokens.labels(trace)
    return (
        parallel_viz.cards(trace),
        gr.update(value=tokens.strip_plain(runs["baseline"], div), label=base_label),
        gr.update(value=tokens.strip_steps(runs["method"], tokens.CHANGED_LABEL[kind]), label=meth_label,
                  color_map={tokens.CHANGED_LABEL[kind]: tokens.COLOR_MAP[tokens.CHANGED_LABEL[kind]]}),
        slider_update(len(steps), first),
        insp,
        gr.update(value=plots.trajectory(trace), visible=kind == "classifier",
                  y_title="P(target)" if trace["summary"].get("attribute") == "sentiment" else "similarity"),
        compact(trace),
        gr.update(value=trace_file(trace), visible=True),
    )


def step_guided(trace, i):
    if not trace:
        return inspector.empty()
    i = int(i or 0)
    return inspector.classifier_step(trace, i) if trace["method"] == "classifier" else inspector.regex_step(trace, i)


def replay_guided(trace):
    if trace:
        yield from replay_steps(trace, tokens.CHANGED_LABEL[trace["method"]])


def guided_method_changed(method, lam, pattern):
    clf = method == "classifier"
    return (
        *(gr.update(visible=clf) for _ in range(5)),  # attribute, lam, k, lookahead, attribute row helpers
        *(gr.update(visible=not clf) for _ in range(3)),  # preset, pattern, top_k
        gr.update(visible=clf),  # trajectory plot
        overview.method_card(method, {"lam": lam}),
    )


def build_guided(family):
    ex = EX.GUIDED[0]
    with gr.Row(equal_height=False):
        with gr.Column(scale=4, min_width=330):
            method = gr.Radio([("Classifier", "classifier"), ("Regex", "regex")], value="classifier", label="Method")
            prompt = gr.Textbox(ex["prompt"], lines=3, label="Prompt")
            with gr.Group() as clf_group:
                attribute = gr.Radio([("sentiment", "sentiment"), ("topic", "topic")], value="sentiment", label="Guide by")
                with gr.Row():
                    target = gr.Radio([("positive", "positive"), ("negative", "negative")], value="positive",
                                      label="Target")
                    topic = gr.Textbox("astronomy", label="Topic", visible=False)
            lam = gr.Slider(0, 50, value=3.0, step=0.5, label="λ · guidance weight")
            k = gr.Slider(2, 40, value=10, step=1, label="k · candidates")
            lookahead = gr.Slider(0, 6, value=0, step=1, label="L · lookahead tokens")
            preset = gr.Dropdown(list(EX.REGEX_PRESETS), value="ISO date", label="Preset", visible=False)
            pattern = gr.Textbox(EX.REGEX_PRESETS["ISO date"], label="Pattern", visible=False)
            top_k = gr.Slider(8, 256, value=64, step=8, label="K · tokens checked", visible=False)
            mode, temperature, max_new, seed, rep, raw = decoding_controls(40, 96)
            run = gr.Button("Run", variant="primary")
        with gr.Column(scale=7, min_width=420):
            note = gr.Markdown(overview.method_card("classifier", {"lam": 3.0}), latex_delimiters=LATEX)
            cards = html()
            base_ht = strip_box("Baseline", tokens.COLORS_BASELINE)
            meth_ht = strip_box("Method", tokens.COLORS_GUIDED, legend=True)
            with gr.Row(equal_height=True):
                step = gr.Slider(0, 1, value=0, step=1, label="Step", scale=6)
                replay = gr.Button("▶ Replay", scale=1, min_width=90)
            insp = html(inspector.empty())
            traj = gr.LinePlot(None, x="step", y="score", color="run", title="Score of the text so far", height=220)
            summary, download = trace_box()

    def fill(e):
        pat = EX.REGEX_PRESETS.get(e.get("preset", "ISO date"), "")
        return [e["method"], e["prompt"], e.get("attribute", "sentiment"), e.get("target", "positive"),
                e.get("topic", "astronomy"), e.get("lam", 3.0), e.get("k", 10), e.get("lookahead", 0),
                e.get("preset", "ISO date"), pat]

    gr.Examples(examples=[fill(e) for e in EX.GUIDED],
                inputs=[method, prompt, attribute, target, topic, lam, k, lookahead, preset, pattern],
                example_labels=[e["label"] for e in EX.GUIDED], label="Examples")
    with gr.Accordion("Papers", open=False):
        html(overview.papers_html("guided"))

    state = gr.State(None)
    outputs = [cards, base_ht, meth_ht, step, insp, traj, summary, download]
    method.change(guided_method_changed, [method, lam, pattern],
                  [clf_group, lam, k, lookahead, attribute, preset, pattern, top_k, traj, note], queue=False)
    attribute.change(lambda a: (gr.update(visible=a == "sentiment"), gr.update(visible=a == "topic"),
                                gr.update(value=3.0 if a == "sentiment" else 20.0)),
                     attribute, [target, topic, lam], queue=False)
    lam.change(lambda m, v: overview.method_card(m, {"lam": v}), [method, lam], note, queue=False)
    preset.change(lambda p, cur: gr.update(value=EX.REGEX_PRESETS.get(p) or cur), [preset, pattern], pattern, queue=False)
    run.click(run_guided,
              [family, method, prompt, attribute, target, topic, lam, k, lookahead, pattern, top_k,
               mode, temperature, max_new, seed, rep, raw],
              state, concurrency_limit=4, concurrency_id="gpu", api_visibility="private"
              ).success(render_guided, state, outputs, queue=False)
    step.change(step_guided, [state, step], insp, queue=False, show_progress="hidden")
    meth_ht.select(select_to_step, None, step, queue=False)
    replay.click(replay_guided, state, meth_ht, concurrency_limit=None, api_visibility="private")
    return state, render_guided, outputs


# ---------------------------------------------------------------------------
# ③ Parallel
# ---------------------------------------------------------------------------


def run_parallel(family, method, prompt, gamma, mode, temperature, num_pred, ngram_max, ngram_min, block, init,
                 max_new, seed, raw):
    P = dict(method=method, prompt=prompt, gamma=gamma, mode=mode, temperature=temperature, num_pred=num_pred,
             ngram_max=ngram_max, ngram_min=ngram_min, block=block, init=init, max_new_tokens=max_new, seed=seed, raw=raw)
    return call_gpu("parallel", family, P)


def first_interesting_iter(trace) -> int:
    for it in trace["runs"]["method"].get("iters", []):
        if any(d[5] == "rej" for d in it["draft"]):
            return it["k"]
    return 0


def render_parallel(trace):
    if not trace:
        return ("", gr.update(value=[]), gr.update(value=[]), slider_update(1, 0), "", "",
                gr.update(value=None), gr.update(), gr.update(), None, gr.update(visible=False))
    runs = trace["runs"]
    s = trace["summary"]
    iters = runs["method"].get("iters", [])
    k = first_interesting_iter(trace)
    spec = trace["method"] == "speculative"
    div = s.get("first_divergence")
    base_label, meth_label, _ = tokens.labels(trace)
    return (
        parallel_viz.cards(trace),
        gr.update(value=tokens.strip_plain(runs["baseline"], div), label=base_label),
        gr.update(value=tokens.strip_kinds(runs["method"]), label=meth_label),
        slider_update(len(iters), k),
        parallel_viz.timeline(trace, k),
        parallel_viz.iteration_detail(trace, k),
        gr.update(value=plots.speedup_curves(s if spec else None)),
        gr.update(value=min(max(s.get("alpha_hat", 0.7), 0.0), 1.0) if spec else 0.7),
        gr.update(value=min(max(s.get("c_hat", 0.5), 0.0), 1.5) if spec else 0.5),
        compact(trace),
        gr.update(value=trace_file(trace), visible=True),
    )


def step_parallel(trace, k):
    if not trace:
        return "", ""
    k = int(k or 0)
    return parallel_viz.timeline(trace, k), parallel_viz.iteration_detail(trace, k)


def select_to_iter(trace, evt: gr.SelectData):
    idx = tokens.select_index(evt)
    if not trace or idx is None:
        return gr.update()
    tok_iter = trace["runs"]["method"].get("tok_iter", [])
    return gr.update(value=tok_iter[idx]) if 0 <= idx < len(tok_iter) else gr.update()


def replay_parallel(trace):
    if not trace:
        return
    run = trace["runs"]["method"]
    value = tokens.strip_kinds(run)
    tok_iter = run.get("tok_iter", [])
    n_iters = len(run.get("iters", []))
    for k in range(n_iters):
        n_tok = sum(1 for t in tok_iter if t <= k)
        yield gr.update(value=value[:n_tok]), parallel_viz.timeline(trace, k, upto=k + 1)
        time.sleep(0.15)
    yield gr.update(value=value), parallel_viz.timeline(trace, None)


def parallel_method_changed(method):
    spec, pld, jac = method == "speculative", method == "pld", method == "jacobi"
    return (
        gr.update(visible=spec), gr.update(visible=spec), gr.update(visible=spec),
        gr.update(visible=pld), gr.update(visible=pld), gr.update(visible=pld),
        gr.update(visible=jac), gr.update(visible=jac),
        gr.update(visible=spec),
        overview.method_card(method, {}),
    )


SPEED_NOTE = (
    "At batch size 1 a forward pass costs about the same whatever the model size; speed tracks the number of layers. "
    "Tokens per target pass is the hardware-independent gain, and the helper's own passes eat into wall-clock time."
)


def build_parallel(family):
    ex = EX.PARALLEL[0]
    with gr.Row(equal_height=False):
        with gr.Column(scale=4, min_width=330):
            method = gr.Radio([("Speculative", "speculative"), ("Prompt lookup", "pld"), ("Jacobi", "jacobi")],
                              value="speculative", label="Method")
            prompt = gr.Textbox(ex["prompt"], lines=3, label="Prompt")
            gamma = gr.Slider(1, 10, value=ex["gamma"], step=1, label="γ · draft length")
            mode = gr.Radio([("greedy", "greedy"), ("sampling", "sampling")], value="greedy", label="Verification")
            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.05, label="Temperature")
            num_pred = gr.Slider(1, 20, value=10, step=1, label="k · tokens copied", visible=False)
            ngram_max = gr.Slider(1, 5, value=3, step=1, label="Max n-gram", visible=False)
            ngram_min = gr.Slider(1, 3, value=1, step=1, label="Min n-gram", visible=False)
            block = gr.Slider(2, 16, value=8, step=1, label="m · block size", visible=False)
            init = gr.Radio([("repeat", "repeat"), ("random", "random")], value="repeat", label="Initial guess",
                            visible=False)
            with gr.Accordion("Settings", open=False):
                max_new = gr.Slider(16, 256, value=96, step=8, label="Max tokens")
                seed = gr.Number(value=0, precision=0, minimum=0, label="Seed")
                raw = gr.Checkbox(value=False, label="Raw prompt (no chat template)")
            run = gr.Button("Run", variant="primary")
        with gr.Column(scale=7, min_width=420):
            note = gr.Markdown(overview.method_card("speculative", ex), latex_delimiters=LATEX)
            cards = html()
            base_ht = strip_box("Baseline", tokens.COLORS_BASELINE)
            meth_ht = strip_box("Method", tokens.COLORS_PARALLEL, legend=True)
            with gr.Row(equal_height=True):
                step = gr.Slider(0, 1, value=0, step=1, label="Target pass", scale=6)
                replay = gr.Button("▶ Replay", scale=1, min_width=90)
            detail = html()
            timeline = html()
            with gr.Accordion("Speedup model", open=False) as whatif:
                gr.Markdown(SPEED_NOTE)
                curve = gr.LinePlot(None, x="γ (draft length)", y="speedup", color="series",
                                    title="Predicted speedup vs γ", height=230)
                with gr.Row():
                    w_alpha = gr.Slider(0.0, 1.0, value=0.7, step=0.01, label="α · acceptance")
                    w_c = gr.Slider(0.0, 1.5, value=0.5, step=0.01, label="c · draft cost")
            summary, download = trace_box()

    def fill(e):
        return [e["method"], e["prompt"], e.get("gamma", 4), e.get("mode", "greedy"), e.get("temperature", 0.8),
                e.get("num_pred", 10), e.get("block", 8)]

    gr.Examples(examples=[fill(e) for e in EX.PARALLEL],
                inputs=[method, prompt, gamma, mode, temperature, num_pred, block],
                example_labels=[e["label"] for e in EX.PARALLEL], label="Examples")
    with gr.Accordion("Papers", open=False):
        html(overview.papers_html("parallel"))

    state = gr.State(None)
    outputs = [cards, base_ht, meth_ht, step, timeline, detail, curve, w_alpha, w_c, summary, download]
    method.change(parallel_method_changed, method,
                  [gamma, mode, temperature, num_pred, ngram_max, ngram_min, block, init, whatif, note], queue=False)
    run.click(run_parallel,
              [family, method, prompt, gamma, mode, temperature, num_pred, ngram_max, ngram_min, block, init, max_new, seed, raw],
              state, concurrency_limit=4, concurrency_id="gpu", api_visibility="private"
              ).success(render_parallel, state, outputs, queue=False)
    step.change(step_parallel, [state, step], [timeline, detail], queue=False, show_progress="hidden")
    meth_ht.select(select_to_iter, state, step, queue=False)
    # .release (not .change): render_parallel sets these sliders to the measured values, and programmatic
    # updates must not redraw the what-if curve with half-updated inputs.
    for s in (w_alpha, w_c):
        s.release(lambda t, a, c: gr.update(value=plots.speedup_curves(t["summary"] if t and t["method"] == "speculative" else None, a, c)),
                  [state, w_alpha, w_c], curve, queue=False, show_progress="hidden")
    replay.click(replay_parallel, state, [meth_ht, timeline], concurrency_limit=None, api_visibility="private")
    return state, render_parallel, outputs


# ---------------------------------------------------------------------------
# Page
# ---------------------------------------------------------------------------


def trace_api(paradigm: str, family: str, params_json: str) -> dict:
    """Run one decoding method plus its baseline and return the full trace."""
    return call_gpu(paradigm, family, json.loads(params_json or "{}"))


with gr.Blocks(title="Beyond Tokens · Decoding Playground", fill_width=True) as demo:
    with gr.Row(equal_height=True):
        with gr.Column(scale=3):
            html(overview.HEADER_HTML)
        with gr.Column(scale=2, min_width=280):
            family = gr.Dropdown(FAMILY_CHOICES, value=DEFAULT_FAMILY, label="Model family")
    if models.LOAD_ERRORS:
        gr.Markdown("<small>Unavailable: " + "; ".join(models.LOAD_ERRORS) + "</small>")

    with gr.Tabs():
        with gr.Tab("Overview"):
            html(overview.FIGURE_SVG)
            html(overview.HOW_TO_HTML)
            html(overview.methods_html())
            with gr.Accordion("📚 Cite the survey", open=True, elem_classes=["dd-cite"]):
                html(f'<div class="dd-wrap dd-cite-text">{overview.CITATION_TEXT}</div>')
                gr.Code(overview.short_bibtex(), language=None, interactive=False, show_label=False,
                        show_line_numbers=False)
            html(overview.LICENSES_HTML)
        with gr.Tab("① Contrastive"):
            c_state, c_render, c_outputs = build_contrastive(family)
        with gr.Tab("② Guided"):
            g_state, g_render, g_outputs = build_guided(family)
        with gr.Tab("③ Parallel"):
            p_state, p_render, p_outputs = build_parallel(family)
        with gr.Tab("Papers"):
            rows = overview.papers_rows()
            gr.Markdown(f"{len(rows)} papers from [Awesome-LLM-Decoding]({overview.GITHUB}).")
            gr.Dataframe(value=rows, headers=["Section", "Paper", "Abbr.", "Venue", "Year", "Model", "Code"],
                         datatype=["str", "markdown", "str", "str", "str", "str", "markdown"], show_search="filter",
                         wrap=True, max_height=680, interactive=False)

    gr.api(trace_api, api_name="trace", api_visibility="undocumented", concurrency_limit=4, concurrency_id="gpu")

    for paradigm, st, render, outs in (("contrastive", c_state, c_render, c_outputs),
                                       ("guided", g_state, g_render, g_outputs),
                                       ("parallel", p_state, p_render, p_outputs)):
        demo.load(lambda p=paradigm: load_gallery(p), None, st, queue=False).then(render, st, outs, queue=False)

demo.queue(max_size=48, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS)