"""Blink v0.5 on ZeroGPU — one forward pass, one calibrated decision per question."""
import os

import spaces  # must come before torch

import html
import json
import math
import time

import gradio as gr
import torch

from blink_core import T_CALIBRATED, Reader, build, load_model, parse_state, render

MODEL_ID = os.environ.get("BLINK_MODEL", "PixilabAI/Blink-v0.5-26B-A4B")
MAXQ = 4
TYPES = {"Choice": "choice", "Yes / No": "noul", "Score": "score"}
LABEL = {v: k for k, v in TYPES.items()}

t0 = time.time()
MODEL, TOK = load_model(MODEL_ID, "cuda")
READER = Reader(MODEL, TOK, "cuda")
print(f"loaded {MODEL_ID} in {time.time() - t0:.0f}s", flush=True)


@spaces.GPU(duration=40, size="xlarge")   # Blink in bf16 is 51.8 GB: over the default half card (48 GB)
def run(state_raw, questions):
    """questions: [(type, question, option lines)] → raw letter log-probs per question (no temperature applied)."""
    state = parse_state(state_raw)
    out, t_all = [], time.perf_counter()
    for kind, question, lines in questions:
        keys, texts = build(kind, lines)
        msgs = render(state, question, texts)
        lps, ms = READER.read(msgs, len(texts))
        out.append({"type": kind, "question": question, "keys": keys, "texts": texts, "logprobs": lps,
                    "ms": round(ms), "prompt": msgs[1]["content"]})
    return {"answers": out, "total_ms": round((time.perf_counter() - t_all) * 1000),
            "structured": not isinstance(state, str)}


# ------------------------------------------------------------------ rendering
def soft(lps, t):
    z = [l / t for l in lps]
    m = max(z)
    e = [math.exp(v - m) for v in z]
    s = sum(e)
    return [v / s for v in e]


def pct(p):
    return "99.9%" if p >= .9995 else "<0.1%" if p < .0005 else f"{p * 100:.1f}%"


EMPTY = '<div class="empty"><div class="lid"></div><p>Nothing to decide yet.</p></div>'


def readout(res, t, animate=True):
    if not res:
        return EMPTY
    t = float(t or T_CALIBRATED)
    n = len(res["answers"])
    parts = [f'<div class="sum"><span><b>{n}</b> decision{"s" if n > 1 else ""}</span>'
             f'<span><b>{res["total_ms"]}</b> ms on the GPU</span><span>{"JSON" if res["structured"] else "text"} state</span>'
             f'<span>T <b>{t:.2f}</b></span></div>']
    for ai, a in enumerate(res["answers"]):
        p = soft(a["logprobs"], t)
        best = max(range(len(p)), key=p.__getitem__)
        bars = []
        for i, (key, text) in enumerate(zip(a["keys"], a["texts"])):
            letter = chr(65 + i)
            k = f"{letter} · {html.escape(key)}" if a["type"] == "choice" and key != text else letter
            bars.append(f'<div class="bar{" win" if i == best else ""}"><div class="bl">'
                        f'<div class="bf{" grow" if animate else ""}" style="width:{p[i] * 100:.2f}%"></div>'
                        f'<div class="bt"><span class="k">{k}</span><span class="d">{html.escape(text)}</span></div></div>'
                        f'<div class="bp">{pct(p[i])}</div></div>')
        extra = ""
        if a["type"] == "noul":
            extra = f'<p class="xs">P(yes) <b>{p[1]:.3f}</b> · answer <b>{"yes" if p[1] >= .5 else "no"}</b></p>'
        elif a["type"] == "score":
            ev = sum(i * v for i, v in enumerate(p))
            extra = (f'<p class="xs">expected level <b>{ev:.2f}</b> on 0–{len(p) - 1} · closest '
                     f'<b>{html.escape(a["texts"][round(ev)])}</b></p>')
        parts.append(
            f'<div class="ans" style="animation-delay:{ai * 70 if animate else 0}ms"><div class="ah">'
            f'<h4>{html.escape(a["question"])}</h4><div class="meta"><span class="tag">{LABEL[a["type"]]}</span>'
            f'<span>{a["ms"]} ms</span></div></div>{"".join(bars)}{extra}'
            f'<details class="prompt"><summary>exact prompt</summary><pre>{html.escape(a["prompt"])}</pre></details></div>')
    return "".join(parts)


from examples import EX  # noqa: E402


# ------------------------------------------------------------------ page
import base64

LOGO = "data:image/png;base64," + base64.b64encode(open(os.path.join(os.path.dirname(__file__), "blink.png"), "rb").read()).decode()
CARD = "https://huggingface.co/PixilabAI/Blink-v0.5-26B-A4B-NVFP4"
HEAD = ('<link rel="preconnect" href="https://fonts.googleapis.com"><link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>'
        '<link href="https://fonts.googleapis.com/css2?family=Instrument+Serif:ital@0;1&family=Geist:wght@300;400;500;600'
        '&family=Geist+Mono:wght@400;500&display=swap" rel="stylesheet">')
HERO = f"""<div class="bk">
<header><a class="brand" href="{CARD}" target="_blank" rel="noopener"><img src="{LOGO}" alt=""><b>Blink</b><span class="pill">v0.5</span></a>
<nav><a class="keep" href="{CARD}" target="_blank" rel="noopener">Model card</a><a href="#how">How it reads</a>
<a href="https://pixilab.ai/?utm_source=hf" target="_blank" rel="noopener">Pixilab</a><a class="keep" href="https://nemini.ai/?utm_source=hf" target="_blank" rel="noopener">Nemini</a></nav></header>
<section class="hero"><div>
<h1>Decide in<br>one <em>blink.</em></h1>
<p class="lede">Blink reads a <strong>state</strong>, a <strong>question</strong> and its <strong>options</strong>, and answers with a
calibrated probability for every option — one forward pass, one token, no prose to parse. Built to sit inside products:
routing, moderation, tagging, gating, dedup, yes/no checks.</p>
<p class="by">Model by <a href="https://pixilab.ai/?utm_source=hf" target="_blank" rel="noopener">pixilab.ai</a> &amp; <a href="https://nemini.ai/?utm_source=hf" target="_blank" rel="noopener">nemini.ai</a></p>
<div class="stats"><div class="stat"><b>58.67</b><span>Decision Index 0.2.1</span></div>
<div class="stat"><b>60.39</b><span>thinking when unsure</span></div>
<div class="stat"><b>0.023</b><span>calibration error</span></div>
<div class="stat"><b>4B</b><span>active of 26B</span></div></div>
</div><div class="eye"><img src="{LOGO}" alt="Blink"></div></section>
<div class="sh"><h2>Try it</h2><p>Pick an example or write your own — up to {MAXQ} questions about one state.</p></div></div>"""
NOTE = ('<p class="note">The slider re-scores the stored answer without touching the GPU. Temperature changes how sure '
        'Blink claims to be — never which option wins. At 1.4 its confidence matches its accuracy.</p>')
HOW = f"""<div class="bk"><div class="sh" id="how"><h2>How it reads</h2><p>The surogate decisions v1 protocol, as on the model card.</p></div>
<section class="how">
<div><b>Lettered options</b><p>Each question becomes one prompt: the state as JSON, the question, the options lettered
<code>A</code>, <code>B</code>, <code>C</code>… Yes/no puts <code>No</code> at A and <code>Yes</code> at B.</p></div>
<div><b>One token</b><p>Thinking off. Blink's answer is the softmax over the option letters' log-probabilities at the
first generated position — nothing is generated beyond it.</p></div>
<div><b>Calibrated</b><p>Divide by <code>T = 1.4</code> before the softmax and the probabilities mean what they say:
threshold on them (<code>P ≥ 0.8</code>) instead of parsing text.</p></div></section>
<div class="bfoot"><span>Blink v0.5 by <a href="https://pixilab.ai/?utm_source=hf" target="_blank" rel="noopener">pixilab.ai</a> &amp; <a href="https://nemini.ai/?utm_source=hf" target="_blank" rel="noopener">nemini.ai</a> · Apache-2.0 · base <a href="https://huggingface.co/google/gemma-4-26B-A4B-it" target="_blank"
rel="noopener">Gemma 4 26B-A4B</a> · this demo runs the BF16 build on ZeroGPU; the scores above are for the
<a href="{CARD}" target="_blank" rel="noopener">NVFP4 build</a> made from it</span>
<span>Not for high-stakes decisions without a threshold and a fallback.</span></div></div>"""
PH = {"Choice": "Options, one per line — key, or key: description",
      "Score": "Levels from lowest to highest, one per line",
      "Yes / No": "Optional — describe the answers:\nfalse: what “no” means\ntrue: what “yes” means"}


def decide(state, n, *fields):
    state = (state or "").strip()
    if not state:
        raise gr.Error("Give Blink a state to look at first.")
    qs = []
    for i in range(int(n)):
        kind, question, options = fields[3 * i: 3 * i + 3]
        if not (question or "").strip():
            continue
        try:
            build(TYPES[kind], (options or "").split("\n"))
        except ValueError as e:
            raise gr.Error(f"Q{i + 1}: {e}.")
        qs.append((TYPES[kind], question.strip(), (options or "").split("\n")))
    if not qs:
        raise gr.Error("Ask at least one question.")
    if len(state) > 24_000:
        raise gr.Error("That state is a bit long for a demo — keep it under 24,000 characters.")
    return run(state, qs)


def decide_api(state: str, questions: list):
    """Score typed decisions about one state with Blink v0.5.

    Args:
        state: the situation — plain text, or a JSON object/array as a string.
        questions: up to 4 items {"type": "choice"|"noul"|"score", "question": str, "options": [str]}.
            choice: "key" or "key: description" per option; score: levels from lowest to highest;
            noul (yes/no): options optional ("false: ...", "true: ...").

    Returns:
        {"answers": [{type, question, keys, probabilities, best, logprobs, ms}], "temperature": 1.4, "total_ms": int}
    """
    if not isinstance(questions, list) or not 1 <= len(questions) <= MAXQ:
        raise gr.Error(f"questions must be a list of 1 to {MAXQ} items")
    qs = []
    for q in questions:
        kind = q.get("type")
        if kind not in ("choice", "noul", "score"):
            raise gr.Error("type must be choice, noul or score")
        try:
            build(kind, list(q.get("options") or []))
        except ValueError as e:
            raise gr.Error(f"{q.get('question', '')[:60]}: {e}")
        qs.append((kind, str(q.get("question", "")).strip(), list(q.get("options") or [])))
    res = run(state, qs)
    for a in res["answers"]:
        p = soft(a["logprobs"], T_CALIBRATED)
        a["probabilities"] = dict(zip(a["keys"], [round(v, 4) for v in p]))
        a["best"] = a["keys"][max(range(len(p)), key=p.__getitem__)]
        a.pop("prompt", None)
    res["temperature"] = T_CALIBRATED
    return res


def load_example(i):
    _, state, qs = EX[i]
    ups = [state, len(qs), None]
    for j in range(MAXQ):
        if j < len(qs):
            k, q, o = qs[j]
            ups += [gr.update(visible=True), gr.update(value=k), q, gr.update(value=o, placeholder=PH[k])]
        else:
            ups += [gr.update(visible=False), gr.update(value="Choice"), "", gr.update(value="", placeholder=PH["Choice"])]
    return ups + [EMPTY]


def show_rows(n):
    return [gr.update(visible=j < n) for j in range(MAXQ)]


with gr.Blocks(title="Blink v0.5 — one forward pass, one calibrated decision") as demo:
    res = gr.State(None)
    nq = gr.State(len(EX[0][2]))
    gr.HTML(HERO)
    with gr.Row(elem_classes="chips"):
        chips = [gr.Button(name, size="sm", scale=0) for name, _, _ in EX]
    with gr.Row(equal_height=False):
        with gr.Column(scale=1, elem_classes="panel"):
            state = gr.Textbox(label="State", value=EX[0][1], lines=7, max_lines=18, elem_id="state",
                               placeholder="What Blink should know — a message, a ticket, a log line, or a JSON object.")
            groups, fields = [], []
            for j in range(MAXQ):
                k0, q0, o0 = EX[0][2][j] if j < len(EX[0][2]) else ("Choice", "", "")
                with gr.Group(visible=j < len(EX[0][2]), elem_classes="q") as g:
                    kind = gr.Radio(list(TYPES), value=k0, show_label=False, container=False, elem_classes="seg")
                    q = gr.Textbox(value=q0, show_label=False, container=False, placeholder=f"Q{j + 1} — one question about the state")
                    o = gr.Textbox(value=o0, show_label=False, container=False, lines=3, max_lines=10,
                                   placeholder=PH[k0], elem_classes="opts")
                    kind.change(lambda k: gr.update(placeholder=PH[k]), kind, o, show_progress="hidden")
                groups.append(g)
                fields += [kind, q, o]
            with gr.Row():
                add = gr.Button("+ Question", size="md", scale=0)
                rm = gr.Button("− Question", size="md", scale=0)
                go = gr.Button("Decide", variant="primary", elem_classes="go", scale=1)
        with gr.Column(scale=1, elem_classes="panel"):
            out = gr.HTML(EMPTY)
            temp = gr.Slider(0.5, 3.0, value=T_CALIBRATED, step=0.05, label="Temperature · 1.4 is calibrated")
            gr.HTML(NOTE)
    gr.HTML(HOW)
    with gr.Row(visible=False):     # API-only endpoint: POST /decide (state, questions) -> JSON
        api_state, api_qs, api_out = gr.Textbox(), gr.JSON(), gr.JSON()
        gr.Button().click(decide_api, [api_state, api_qs], api_out, api_name="decide")

    go.click(decide, [state, nq, *fields], res).then(lambda r, t: readout(r, t, True), [res, temp], out)
    temp.change(lambda r, t: readout(r, t, False), [res, temp], out, show_progress="hidden")
    add.click(lambda n: min(n + 1, MAXQ), nq, nq).then(show_rows, nq, groups, show_progress="hidden")
    rm.click(lambda n: max(n - 1, 1), nq, nq).then(show_rows, nq, groups, show_progress="hidden")
    row_outputs = [state, nq, res]
    for j in range(MAXQ):
        row_outputs += [groups[j], *fields[3 * j: 3 * j + 3]]
    for i, b in enumerate(chips):
        b.click(lambda i=i: load_example(i), None, row_outputs + [out], show_progress="hidden")

if __name__ == "__main__":
    demo.queue(max_size=40).launch(css=open(os.path.join(os.path.dirname(__file__), "style.css")).read(), head=HEAD)