import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch)

import html
import json
import random
import re
import time

import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from safetensors import safe_open
from transformers import AutoConfig, AutoProcessor
from transformers.activations import ACT2FN
from transformers.cache_utils import DynamicCache
from transformers.initialization import no_init_weights
from transformers.models.diffusion_gemma import modeling_diffusion_gemma as mdg

MODEL_ID = "Singularity-Compute/super-decider"

# ---------------------------------------------------------------------------------------------
# Server settings, matching runtime/Dockerfile of github.com/Singularity-Compute/super-decider
# ---------------------------------------------------------------------------------------------
CANVAS_LEN = 128  # CANVAS=128
CANVAS_STEP = 16
VOCAB = 262144
TURN_CLOSE = 106
PAD = 0
TEMPS = {"score": 0.375}  # DJEV_TEMPS=score=0.375
NOUL_COMMIT = 0.8  # DJEV_NOUL_COMMIT=0.8
ADAPT_GATES = {"choice": 0.7, "score": 0.7, "noul": 0.85}  # DJEV_ADAPTIVE_THINK
ADAPT_BUDGET = 512
ADAPT_MAX_Q = 2
THINK_STEPS = 8  # DJEV_THINK_STEPS=8
BIDIR_MAXPROMPT = 8192  # DJEV_BIDIR_MAXPROMPT
MAX_QUESTIONS = 64
MAX_PROMPT_TOKENS = 32768  # demo cap
MAX_IMAGES = 8
SCAFFOLD_TEXT = "<|channel>thought\n<channel|>"

# ---------------------------------------------------------------------------------------------
# Model: DiffusionGemma 26B-A4B with ModelOpt NVFP4 experts (W4A16, dequantized on the fly on GPU)
# ---------------------------------------------------------------------------------------------
_FP4 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]
_LUT2 = torch.tensor([[_FP4[b & 15], _FP4[b >> 4]] for b in range(256)], dtype=torch.bfloat16)


def _dequant(q, s, g, lut2):
    """q [n, R, C/2] uint8 (low nibble first), s [n, R, C/16] fp8 block scales, g [n, R, 1] fp32 global."""
    n, r, c2 = q.shape
    v = lut2[q.int()].view(n, r, c2 // 8, 16)
    sc = s.float() * g
    return (v.float() * sc[..., None]).to(torch.bfloat16).view(n, r, c2 * 2)


class NVFP4Experts(nn.Module):
    """Drop-in for DiffusionGemmaTextExperts holding NVFP4-packed weights."""

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_experts
        self.act_fn = ACT2FN[config.hidden_activation]
        # tiny placeholders so the stock _init_weights isinstance branch is harmless; removed after build
        self.gate_up_proj = nn.Parameter(torch.zeros(1))
        self.down_proj = nn.Parameter(torch.zeros(1))
        self.register_buffer("lut2", _LUT2.clone(), persistent=False)

    def forward(self, hidden_states, top_k_index, top_k_weights):
        out = torch.zeros_like(hidden_states)
        with torch.no_grad():
            mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)  # [E, K, T]
            hit = torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero().flatten()
        if hit.numel() == 0:
            return out
        gu = _dequant(self.gu_q[hit], self.gu_s[hit], self.gu_g[hit], self.lut2)
        dn = _dequant(self.dn_q[hit], self.dn_s[hit], self.dn_g[hit], self.lut2)
        for j, e in enumerate(hit.tolist()):
            k_pos, tok = torch.where(mask[e])
            x = hidden_states[tok]
            gate, up = F.linear(x, gu[j]).chunk(2, dim=-1)
            h = F.linear(self.act_fn(gate) * up, dn[j])
            h = h * top_k_weights[tok, k_pos, None]
            out.index_add_(0, tok, h.to(out.dtype))
        return out


def _set_param(root, path, value):
    *mods, leaf = path.split(".")
    m = root
    for p in mods:
        m = getattr(m, p)
    setattr(m, leaf, value)


def load_model():
    local = snapshot_download(MODEL_ID)
    config = AutoConfig.from_pretrained(local)
    if hasattr(config, "quantization_config"):
        del config.quantization_config
    config.canvas_length = CANVAS_LEN  # the served canvas (thought generation)
    tc = config.text_config
    E, H, I = tc.num_experts, tc.hidden_size, tc.moe_intermediate_size
    L = tc.num_hidden_layers

    orig = mdg.DiffusionGemmaTextExperts
    mdg.DiffusionGemmaTextExperts = NVFP4Experts
    try:
        with no_init_weights():
            model = mdg.DiffusionGemmaForBlockDiffusion._from_config(
                config, dtype=torch.bfloat16, attn_implementation="sdpa"
            )
    finally:
        mdg.DiffusionGemmaTextExperts = orig
    # buffers that are normally computed by _init_weights (rotary tables, embed scales)
    for m in model.modules():
        if isinstance(m, (mdg.DiffusionGemmaTextRotaryEmbedding, mdg.DiffusionGemmaTextScaledWordEmbedding)):
            model._init_weights(m)

    idx = json.load(open(os.path.join(local, "model.safetensors.index.json")))["weight_map"]
    files = {f: safe_open(os.path.join(local, f), framework="pt") for f in set(idx.values())}

    def get(k):
        return files[idx[k]].get_tensor(k)

    dec = model.model.decoder
    t0 = time.time()
    for li in range(L):
        ex = dec.layers[li].experts
        gu_q = torch.empty(E, 2 * I, H // 2, dtype=torch.uint8)
        gu_s = torch.empty(E, 2 * I, H // 16, dtype=torch.float8_e4m3fn)
        gu_g = torch.empty(E, 2 * I, 1, dtype=torch.float32)
        dn_q = torch.empty(E, H, I // 2, dtype=torch.uint8)
        dn_s = torch.empty(E, H, I // 16, dtype=torch.float8_e4m3fn)
        dn_g = torch.empty(E, H, 1, dtype=torch.float32)
        for e in range(E):
            p = f"model.decoder.layers.{li}.experts.{e}."
            gu_q[e, :I] = get(p + "gate_proj.weight")
            gu_q[e, I:] = get(p + "up_proj.weight")
            gu_s[e, :I] = get(p + "gate_proj.weight_scale")
            gu_s[e, I:] = get(p + "up_proj.weight_scale")
            gu_g[e, :I] = get(p + "gate_proj.weight_scale_2").float()
            gu_g[e, I:] = get(p + "up_proj.weight_scale_2").float()
            dn_q[e] = get(p + "down_proj.weight")
            dn_s[e] = get(p + "down_proj.weight_scale")
            dn_g[e] = get(p + "down_proj.weight_scale_2").float()
        del ex.gate_up_proj, ex.down_proj
        ex.register_buffer("gu_q", gu_q)
        ex.register_buffer("gu_s", gu_s)
        ex.register_buffer("gu_g", gu_g)
        ex.register_buffer("dn_q", dn_q)
        ex.register_buffer("dn_s", dn_s)
        ex.register_buffer("dn_g", dn_g)
    print(f"experts packed in {time.time() - t0:.0f}s", flush=True)

    sd = {k: get(k) for k in idx if ".experts." not in k}
    res = model.load_state_dict(sd, strict=False)
    del sd

    # tie the encoder text stack (and lm_head) to the decoder, as in the checkpoint
    enc = model.model.encoder.language_model
    tied = 0
    for el, dl in zip(enc.layers, dec.layers):
        el.experts = dl.experts
        for name, p in dl.named_parameters():
            _set_param(el, name, p)
            tied += 1
    enc.norm.weight = dec.norm.weight
    enc.embed_tokens.weight = dec.embed_tokens.weight
    model.lm_head.weight = dec.embed_tokens.weight
    missing = [
        k for k in res.missing_keys
        if not (k.startswith("model.encoder.language_model.") or k == "lm_head.weight" or ".experts." in k)
    ]
    print(f"tied {tied} encoder params; unexpected={res.unexpected_keys[:5]} missing={missing[:10]}", flush=True)
    files.clear()
    return local, model.eval()


LOCAL_DIR, model = load_model()
model.to("cuda")
processor = AutoProcessor.from_pretrained(LOCAL_DIR)
TOK = processor.tokenizer
EOS = int(TOK.eos_token_id) if TOK.eos_token_id is not None else 1
THOUGHT_OPEN = TOK.encode("<|channel>thought\n", add_special_tokens=False)
THOUGHT_CLOSE = TOK.encode("<channel|>", add_special_tokens=False)
SCAFFOLD = TOK.encode(SCAFFOLD_TEXT, add_special_tokens=False)
assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, "the thought tags must tokenize apart"
SLIDING_WINDOW = model.config.text_config.sliding_window
SOFTCAP = model.config.text_config.final_logit_softcapping


class SchemaError(ValueError):
    pass


# ---------------------------------------------------------------------------------------------
# Schema, prompt and answer template (ported from runtime/structured_server.py)
# ---------------------------------------------------------------------------------------------
FORMATS = {
    "lines": ("\n", "{id}: ", 'Reply with one line per question, in this order, formatted as "id: label".'),
    "indexed": (
        " ",
        "{id}",
        "Reply on one line with each question's id immediately followed by its label, separated by single spaces.",
    ),
}
DECISION_PREAMBLE = (
    "Answer a fixed set of questions about the state the user provides. "
    "Each question lists its allowed answers; reply with exactly one label "
    "per question.\n"
)


def enc(text):
    return TOK.encode(text, add_special_tokens=False)


def jev_schema(body):
    qs = body.get("questions")
    if not isinstance(qs, dict) or not qs:
        raise SchemaError("questions: needs a non-empty map of id -> question")
    if len(qs) > MAX_QUESTIONS:
        raise SchemaError(f"at most {MAX_QUESTIONS} questions")
    out = []
    for qid, q in qs.items():
        qid = str(qid).strip()
        if not qid or ":" in qid or "\n" in qid:
            raise SchemaError(f"question id {qid!r} must be non-empty, no ':' or newline")
        if not isinstance(q, dict):
            raise SchemaError(f"question {qid!r}: must be an object")
        if any(k in q for k in ("depends_on", "ask_if")):
            raise SchemaError(f"question {qid!r}: depends_on / ask_if are not supported in this demo")
        kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "")
        ins = ins if isinstance(ins, str) else json.dumps(ins)
        if kind == "noul":
            if crit is not None and not isinstance(crit, dict):
                raise SchemaError(f"question {qid!r}: noul criteria must be an object with true and false")
            crit = crit or {}
            choices = [("yes", crit.get("true")), ("no", crit.get("false"))]
            labels = ["yes", "no"]
        elif kind == "choice":
            if not isinstance(crit, dict) or not crit:
                raise SchemaError(f"question {qid!r}: choice criteria must map option names to descriptions")
            choices = [(str(n), d) for n, d in crit.items()]
            labels = [chr(ord("A") + i) for i in range(len(choices))]
        elif kind == "score":
            if not isinstance(crit, list):
                raise SchemaError(f"question {qid!r}: score criteria must be an ordered list of levels")
            choices = [(str(level), None) for level in crit]
            labels = (
                [str(i + 1) for i in range(len(choices))]
                if len(choices) <= 9
                else [chr(ord("A") + i) for i in range(len(choices))]
            )
        elif kind in ("span", "spans"):
            raise SchemaError(f"question {qid!r}: span questions are not supported in this demo")
        else:
            raise SchemaError(f"question {qid!r}: unknown type {kind!r} (use noul, choice or score)")
        if len(choices) < 2:
            raise SchemaError(f"question {qid!r}: needs at least two alternatives")
        if len(choices) > 26:
            raise SchemaError(f"question {qid!r}: at most 26 alternatives in this demo")
        out.append({"id": qid, "type": kind, "instructions": ins, "choices": choices, "labels": labels})
    fmt = "lines" if len(out) <= 10 else "indexed"
    schema = {"questions": out, "instructions": body.get("instructions"), "format": fmt}
    if fmt == "indexed" and not all(fits_indexed(q) for q in out):
        schema["format"] = "lines"
    return schema


def fits_indexed(q):
    try:
        resolve_template([q], [], "", "indexed")
        return True
    except SchemaError:
        return False


def system_text(schema):
    s = DECISION_PREAMBLE
    if schema.get("instructions"):
        s += "\n" + str(schema["instructions"]).strip() + "\n"
    for q in schema["questions"]:
        s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n"
        for (name, desc), label in zip(q["choices"], q["labels"]):
            if q["type"] == "noul":
                s += f"  {label}: {str(desc).strip()}\n" if desc else f"  {label}\n"
            elif desc:
                s += f"  {label}: {name} ({str(desc).strip()})\n"
            else:
                s += f"  {label}: {name}\n"
    s += "\n" + FORMATS[schema.get("format", "lines")][2]
    return s


def chat_messages(sys_text, state_content):
    """DJEV_DOC_FIRST=1: the fixed preamble is the system turn; the questions follow the state."""
    rest = sys_text[len(DECISION_PREAMBLE):].lstrip("\n")
    if isinstance(state_content, str):
        user = state_content + "\n\n" + rest
    else:
        user = list(state_content) + [{"type": "text", "text": "\n\n" + rest}]
    return [{"role": "system", "content": DECISION_PREAMBLE.rstrip("\n")}, {"role": "user", "content": user}], rest


def answer_text(qs, labels, fmt="lines"):
    join, lead, _ = FORMATS[fmt]
    return join.join(lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels))


def resolve_template(qs, head, lead, fmt):
    base_labels = [0] * len(qs)
    base = head + enc(lead + answer_text(qs, base_labels, fmt))
    if len(base) + 1 > CANVAS_LEN:
        raise SchemaError(f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}")
    slots = []
    for qi, q in enumerate(qs):
        pos = None
        ids = [0] * len(q["labels"])
        for li in range(1, len(q["labels"])):
            labels = list(base_labels)
            labels[qi] = li
            e = head + enc(lead + answer_text(qs, labels, fmt))
            if len(e) != len(base):
                raise SchemaError(f"question {q['id']!r}: label {q['labels'][li]!r} is not a single token")
            diffs = [i for i in range(len(e)) if e[i] != base[i]]
            if len(diffs) != 1 or (pos is not None and diffs[0] != pos):
                raise SchemaError(f"question {q['id']!r}: labels do not share one template slot")
            pos = diffs[0]
            ids[li] = e[pos]
        ids[0] = base[pos]
        if len(set(ids)) != len(ids):
            raise SchemaError(f"question {q['id']!r}: two labels tokenize to the same id")
        slots.append({"pos": pos, "label_ids": ids})
    return base, slots


def chunk_groups(schema):
    groups, group = [], []
    for q in schema["questions"]:
        trial = group + [q]
        rows = len(SCAFFOLD) + len(enc(answer_text(trial, [0] * len(trial), schema["format"]))) + 1
        if rows > CANVAS_LEN and group:
            groups.append(group)
            group = [q]
        else:
            group = trial
    if group:
        groups.append(group)
    return groups


def canvas_width(template):
    need = len(template) + 1
    return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP)


def build_canvas(template, slots, seed):
    rng = random.Random(seed)
    canvas = list(template) + [TURN_CLOSE]
    canvas += [PAD] * (canvas_width(template) - len(canvas))
    for s in slots:
        canvas[s["pos"]] = rng.randrange(VOCAB)
    return canvas


def bidir_ranges(text, user, rest_len):
    """DJEV_BIDIR=qsep: one bidirectional range per question block (question line + its options)."""
    e_ = TOK(text, add_special_tokens=False, return_offsets_mapping=True)
    ids, offs = e_["input_ids"], e_["offset_mapping"]
    if len(ids) > BIDIR_MAXPROMPT:
        return ids, None
    e = text.rfind(user) + len(user)
    a = e - rest_len
    rest = user[-rest_len:]
    starts = [m.start() for m in re.finditer(r"(?m)^Question \S+: ", rest)]
    endq = rest.rfind("\n\nReply with")
    endq = endq if endq > 0 else len(rest)
    out = []
    for k, s0 in enumerate(starts):
        c0, c1 = a + s0, a + (starts[k + 1] - 2 if k + 1 < len(starts) else endq)
        jdx = [i for i, (x, y) in enumerate(offs) if x >= c0 and y <= c1 and y > x]
        if jdx and jdx[-1] - jdx[0] + 1 == len(jdx) and jdx[-1] > jdx[0]:
            out.append((jdx[0], jdx[-1] + 1))
    return ids, (out or None)


def text_prompt(sys_text, state_text, thinking):
    msgs, rest = chat_messages(sys_text, state_text)
    text = TOK.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=thinking)
    ids, ranges = bidir_ranges(text, msgs[1]["content"], len(rest))
    return ids, ranges


def image_prompt(sys_text, state_parts, thinking):
    msgs, _ = chat_messages(sys_text, state_parts)
    content = []
    for p in msgs[1]["content"]:
        if p["type"] == "image":
            content.append({"type": "image", "image": p["image"]})
        else:
            content.append({"type": "text", "text": p["text"]})
    hf_msgs = [{"role": "system", "content": [{"type": "text", "text": msgs[0]["content"]}]},
               {"role": "user", "content": content}]
    inputs = processor.apply_chat_template(
        hf_msgs, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt",
        enable_thinking=thinking,
    )
    return inputs


# ---------------------------------------------------------------------------------------------
# GPU reads
# ---------------------------------------------------------------------------------------------
def encoder_masks(n, ranges, device):
    i = torch.arange(n, device=device)
    causal = i[None, :] <= i[:, None]
    window = (i[:, None] - i[None, :]) < SLIDING_WINDOW
    full, slide = causal, causal & window
    if ranges:
        rid = torch.full((n,), -1, device=device, dtype=torch.long)
        for k, (a, b) in enumerate(ranges):
            rid[a:b] = k
        same = (rid[:, None] == rid[None, :]) & (rid[:, None] >= 0)
        full, slide = full | same, slide | same
    return {"full_attention": full[None, None], "sliding_attention": slide[None, None]}


def encode_prompt(prompt_ids, ranges=None, mm=None, upto=None):
    """Prefill the encoder KV cache. Text prompts get the per-question bidirectional mask."""
    upto = len(prompt_ids) if upto is None else upto
    ids = torch.tensor([prompt_ids[:upto]], device="cuda")
    cache = DynamicCache(config=model.config.get_text_config(decoder=True))
    pos = torch.arange(upto, device="cuda")[None]
    kw = {}
    if mm is not None:
        kw = {k: v.to("cuda") for k, v in mm.items() if k in ("pixel_values", "image_position_ids", "mm_token_type_ids")}
        if "pixel_values" in kw:
            kw["pixel_values"] = kw["pixel_values"].to(torch.bfloat16)
        masks = model.model.encoder.create_masks_for_generate(
            config=model.config,
            inputs_embeds=torch.empty((1, upto, 0), dtype=torch.bfloat16, device="cuda"),
            attention_mask=torch.ones((1, upto), dtype=torch.bool, device="cuda"),
            past_key_values=cache,
            position_ids=pos,
            mm_token_type_ids=kw.get("mm_token_type_ids"),
        )
    else:
        masks = encoder_masks(upto, [r for r in (ranges or []) if r[1] <= upto], "cuda")
    model.model.encoder(input_ids=ids, attention_mask=masks, past_key_values=cache, position_ids=pos, **kw)
    return cache


@torch.no_grad()
def read(prompt_ids, ranges, template, slots, seed, mm=None):
    """One denoise step over a seeded canvas: label probabilities at every answer slot."""
    cache = encode_prompt(prompt_ids, ranges, mm)
    n = len(prompt_ids)
    canvas = build_canvas(template, slots, seed)
    w = len(canvas)
    out = model.model.decoder(
        decoder_input_ids=torch.tensor([canvas], device="cuda"),
        past_key_values=cache,
        decoder_position_ids=torch.arange(n, n + w, device="cuda")[None],
    )
    hs = out.last_hidden_state[0, [s["pos"] for s in slots]]
    logits = model.lm_head(hs).float()
    logits = torch.tanh(logits / SOFTCAP) * SOFTCAP
    res = []
    for j, s in enumerate(slots):
        lp = logits[j, s["label_ids"]]
        res.append(torch.softmax(lp, -1).cpu().tolist())
    return res


@torch.no_grad()
def write_thought(prompt_ids, ranges, budget, mm=None, seed=42):
    """Generate a thought (block diffusion, DJEV_THINK_STEPS denoising steps) until the close tag."""
    torch.manual_seed(seed)
    close = THOUGHT_CLOSE[0]
    gen_kw = dict(
        max_new_tokens=int(budget), max_denoising_steps=THINK_STEPS, eos_token_id=[1, 106, 50, close],
        return_dict_in_generate=True,
    )
    if mm is None:
        cache = encode_prompt(prompt_ids, ranges, None, upto=len(prompt_ids) - 1)
        out = model.generate(
            input_ids=torch.tensor([prompt_ids[-1:]], device="cuda"), past_key_values=cache, **gen_kw
        )
        new = out.sequences[0, 1:].tolist()
    else:
        kw = {k: v.to("cuda") for k, v in mm.items() if k != "attention_mask"}
        if "pixel_values" in kw:
            kw["pixel_values"] = kw["pixel_values"].to(torch.bfloat16)
        out = model.generate(**kw, **gen_kw)
        new = out.sequences[0, mm["input_ids"].shape[1]:].tolist()
    if new[: len(THOUGHT_OPEN)] == THOUGHT_OPEN:
        new = new[len(THOUGHT_OPEN):]
    cut = len(new)
    for t in (close, 1, 106, 50):
        if t in new:
            cut = min(cut, new.index(t))
    return new[:cut]


def summarize(q, probs):
    mean = list(probs)
    raw_top = max(mean)
    t = TEMPS.get(q["type"], 1.0)
    if t != 1.0:
        w = [max(m, 1e-12) ** (1.0 / t) for m in mean]
        mean = [x / sum(w) for x in w]
    if NOUL_COMMIT and q["type"] == "noul" and 1 - NOUL_COMMIT < mean[0] < NOUL_COMMIT:
        y = NOUL_COMMIT + 0.001 if mean[0] >= 0.5 else 1 - NOUL_COMMIT - 0.001
        mean = [y, 1 - y]
    top = max(range(len(mean)), key=lambda i: mean[i])
    names = [c[0] for c in q["choices"]]
    if q["type"] == "noul":
        ans = {"type": "noul", "noul": mean[0]}
    elif q["type"] == "choice":
        ans = {"type": "choice", "choice": names[top], "probabilities": dict(zip(names, mean)),
               "confidence": mean[top]}
    else:
        ans = {
            "type": "score",
            "score": sum(i * p for i, p in enumerate(mean)),
            "legend": {str(i): n for i, n in enumerate(names)},
            "probabilities": {str(i): p for i, p in enumerate(mean)},
            "confidence": mean[top],
        }
    return ans, raw_top


def decide_group(sub, state, seed, think_budget, images_mm_fn):
    """One joint read of a question group, optionally preceded by a thought."""
    sys_text = system_text(sub)
    thought_ids = None
    if isinstance(state, str):
        if think_budget:
            ids, ranges = text_prompt(sys_text, state, thinking=True)
            prompt = ids + THOUGHT_OPEN
            thought_ids = write_thought(prompt, ranges, think_budget, seed=seed)
            prompt = prompt + thought_ids + THOUGHT_CLOSE
            template, slots = resolve_template(sub["questions"], [], "", sub["format"])
        else:
            prompt, ranges = text_prompt(sys_text, state, thinking=False)
            template, slots = resolve_template(sub["questions"], SCAFFOLD, "", sub["format"])
        if len(prompt) > MAX_PROMPT_TOKENS:
            raise SchemaError(f"state too long for this demo ({len(prompt)} tokens > {MAX_PROMPT_TOKENS})")
        probs = read(prompt, ranges, template, slots, seed)
        n_prompt = len(prompt)
    else:
        head = SCAFFOLD
        thinking = False
        if think_budget:
            answer = enc(answer_text(sub["questions"], [0] * len(sub["questions"]), sub["format"]))
            fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer)
            if fits >= 8:
                mm_t = images_mm_fn(sys_text, True)
                thought_ids = write_thought(None, None, min(think_budget, fits), mm=mm_t, seed=seed)
                head = THOUGHT_OPEN + thought_ids + THOUGHT_CLOSE
                thinking = True
        mm = images_mm_fn(sys_text, thinking)
        template, slots = resolve_template(sub["questions"], head, "", sub["format"])
        prompt = mm["input_ids"][0].tolist()
        probs = read(prompt, None, template, slots, seed, mm=mm)
        n_prompt = len(prompt)
    return probs, n_prompt, (len(template) + 1 + (len(thought_ids) if thought_ids else 0)), thought_ids


def run_decision(body, image_paths, thinking_mode, seed):
    schema = jev_schema(body)
    state = body.get("state")
    if state is None or state == "":
        raise SchemaError("state: required")
    text = state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)
    if image_paths:
        from PIL import Image

        imgs = [Image.open(p).convert("RGB") for p in image_paths[:MAX_IMAGES]]
        state_content = [{"type": "image", "image": im} for im in imgs] + [{"type": "text", "text": text}]
    else:
        state_content = text

    def mm_fn(sys_text, thinking):
        return image_prompt(sys_text, state_content, thinking)

    always = thinking_mode == "always"
    budget = ADAPT_BUDGET if always else 0
    answers, raw, diag = {}, {}, {"reads": 0, "thought_tokens": 0, "adaptive_think": []}
    in_tok, out_tok = 0, 0
    groups = chunk_groups(schema)
    for k, group in enumerate(groups):
        sub = dict(schema, questions=group)
        gseed = seed + 104729 * k
        probs, n_prompt, n_out, th = decide_group(sub, state_content, gseed, budget, mm_fn)
        diag["reads"] += 1
        in_tok = max(in_tok, n_prompt)
        out_tok += n_out
        if th:
            diag["thought_tokens"] += len(th)
        for q, p in zip(group, probs):
            answers[q["id"]], raw[q["id"]] = summarize(q, p)
        # adaptive thinking: re-answer uncertain questions after a thought (text states, <= 2 questions)
        if thinking_mode == "auto" and isinstance(state_content, str) and len(group) <= ADAPT_MAX_Q:
            low = [q for q in group if raw[q["id"]] < ADAPT_GATES.get(q["type"], 0.0)]
            if low:
                try:
                    probs2, _, n2, th2 = decide_group(sub, state_content, gseed, ADAPT_BUDGET, mm_fn)
                except SchemaError:
                    continue
                diag["reads"] += 1
                out_tok += n2
                diag["thought_tokens"] += len(th2 or [])
                diag["adaptive_think"] += [q["id"] for q in low]
                for q, p in zip(group, probs2):
                    if q in low:
                        answers[q["id"]], _ = summarize(q, p)
    return schema, answers, {"input_tokens": in_tok, "output_tokens": out_tok}, diag


# ---------------------------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------------------------
def render(schema, answers):
    cards = []
    for q in schema["questions"]:
        a = answers[q["id"]]
        names = [c[0] for c in q["choices"]]
        if a["type"] == "noul":
            rows = [("yes", a["noul"]), ("no", 1 - a["noul"])]
            headline = f"P(yes) = {a['noul']:.2f}"
            best = "yes" if a["noul"] >= 0.5 else "no"
        elif a["type"] == "choice":
            rows = [(n, a["probabilities"][n]) for n in names]
            headline = f"{html.escape(a['choice'])} &middot; {a['confidence']:.2f}"
            best = a["choice"]
        else:
            rows = [(n, a["probabilities"][str(i)]) for i, n in enumerate(names)]
            best = names[max(range(len(names)), key=lambda i: a["probabilities"][str(i)])]
            headline = f"score {a['score']:.2f} (0&ndash;{len(names) - 1}) &middot; {html.escape(best)}"
        bars = "".join(
            f"<div class='sd-row'><span class='sd-name{' sd-best' if n == best else ''}'>{html.escape(n)}</span>"
            f"<span class='sd-bar'><span style='width:{max(p, 0.003) * 100:.1f}%'></span></span>"
            f"<span class='sd-p'>{p:.2f}</span></div>"
            for n, p in rows
        )
        cards.append(
            f"<div class='sd-card'><div class='sd-head'><b>{html.escape(q['id'])}</b>"
            f"<span class='sd-type'>{a['type']}</span><span class='sd-hl'>{headline}</span></div>"
            f"<div class='sd-q'>{html.escape(q['instructions'])}</div>{bars}</div>"
        )
    return "<div class='sd-wrap'>" + "".join(cards) + "</div>"


def _image_paths(images):
    paths = []
    for it in images or []:
        if isinstance(it, (list, tuple)):
            it = it[0]
        if isinstance(it, dict):
            it = it.get("path") or it.get("name")
        if it:
            paths.append(it)
    return paths


def _gpu_duration(state, questions, thinking="auto", images=None, seed=42, *args, **kwargs):
    """Measured on ZeroGPU: one read ~2.5 s, a 512-token thought + read ~24 s."""
    try:
        qs = json.loads(questions) if isinstance(questions, str) else questions
        if isinstance(qs, dict) and isinstance(qs.get("questions"), dict):
            qs = qs["questions"]
        n = len(qs)
    except Exception:
        n = 1
    if thinking == "always" or (thinking == "auto" and n <= ADAPT_MAX_Q and not images):
        return 30
    return 10 if n <= 16 else 20


@spaces.GPU(duration=_gpu_duration)
def decide(state: str, questions: str, thinking: str = "auto", images: list | None = None, seed: int = 42):
    """Answer typed decision questions about a state with calibrated probabilities.

    Args:
        state: The state to decide about: plain text, or JSON (an object or array).
        questions: Jev / System One questions as JSON: a map of id -> {"type": "noul"|"choice"|"score",
            "instructions": str, "criteria": ...}. choice criteria map option names to descriptions,
            score criteria are an ordered list of levels, noul criteria are optional {"true", "false"}.
        thinking: "off", "auto" (a short thought only on uncertain questions) or "always".
        images: Optional list of image file paths (up to 8) shown ahead of the state.
        seed: Seed for the answer-slot noise of the canvas.

    Returns:
        An HTML summary and the Jev-shaped JSON answer set.
    """
    t0 = time.time()
    try:
        qs = json.loads(questions) if isinstance(questions, str) else questions
        if isinstance(qs, dict) and "questions" in qs and isinstance(qs["questions"], dict):
            qs = qs["questions"]
    except json.JSONDecodeError as e:
        raise gr.Error(f"Questions must be valid JSON: {e}")
    st = state
    if isinstance(state, str) and state.strip()[:1] in ("{", "["):
        try:
            st = json.loads(state)
        except json.JSONDecodeError:
            st = state
    try:
        schema, answers, usage, diag = run_decision(
            {"state": st, "questions": qs}, _image_paths(images), thinking or "auto", int(seed)
        )
    except SchemaError as e:
        raise gr.Error(str(e))
    elapsed = time.time() - t0
    result = {
        "model": "super-decider",
        "answers": answers,
        "usage": usage,
        "diagnostics": {**diag, "seconds": round(elapsed, 3)},
    }
    print(f"decide: {len(answers)} q, reads={diag['reads']} thought={diag['thought_tokens']} {elapsed:.2f}s", flush=True)
    return render(schema, answers), result


EX_TICKET_STATE = (
    'Order 8812 shipped on Oct 1. The customer writes: "I want to cancel and get my money back." '
    "Policy: shipped orders cannot be cancelled; they can be returned within 30 days of delivery for a full refund."
)
EX_TICKET_Q = {
    "next_step": {"type": "choice", "instructions": "What should support do next?",
                  "criteria": {"refund_now": "Refund without a return", "start_return": "Start a return",
                               "escalate": "Hand over to a human"}},
    "allowed": {"type": "noul", "instructions": "Is cancelling this order allowed by the policy?"},
    "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "medium", "high"]},
}
EX_INVOICE_STATE = "Invoice total 1,240 EUR. The customer paid 2 x 400 EUR and 1 x 340 EUR."
EX_INVOICE_Q = {
    "paid": {"type": "noul", "instructions": "Is the invoice fully paid?"},
    "next_step": {"type": "choice", "instructions": "What should accounts receivable do next?",
                  "criteria": {"close": "Mark as paid and close", "remind": "Send a reminder for the remaining amount",
                               "escalate": "Hand over to collections"}},
}
EX_ROUTE_STATE = {"ticket": "I was charged twice for my March invoice.", "plan": "Pro", "customer_since": "2024-02"}
EX_ROUTE_Q = {
    "team": {"type": "choice", "instructions": "Which team should handle this ticket?",
             "criteria": {"billing": "Payments, invoices, refunds", "technical": "Bugs and outages",
                          "account": "Login, profile, plan changes", "sales": "New plans and upgrades"}},
    "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "normal", "high"]},
    "churn_risk": {"type": "noul", "instructions": "Does the customer sound likely to cancel?"},
}
EX_IMAGE_STATE = "Scanned document from the accounts payable inbox."
EX_IMAGE_Q = {
    "doc_type": {"type": "choice", "instructions": "What kind of document is this?",
                 "criteria": {"invoice": "An invoice", "receipt": "A store receipt", "contract": "A contract",
                              "other": "Something else"}},
    "biggest_line": {"type": "choice", "instructions": "Which line item has the largest amount?",
                     "criteria": {"paper": "Paper A4 (box)", "toner": "Toner cartridge", "lamp": "Desk lamp"}},
    "total_ok": {"type": "noul", "instructions": "Does the TOTAL equal the subtotal plus tax?"},
    "vendor": {"type": "choice", "instructions": "Who issued the document?",
               "criteria": {"acme": "ACME Office Supplies", "northwind": "Northwind Ltd", "unknown": "Not shown"}},
}


def _j(x):
    return json.dumps(x, indent=2, ensure_ascii=False)


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.sd-wrap { display: flex; flex-direction: column; gap: 12px; }
.sd-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 10px 14px; }
.sd-head { display: flex; gap: 10px; align-items: baseline; flex-wrap: wrap; }
.sd-type { font-size: 0.8em; opacity: 0.7; border: 1px solid var(--border-color-primary); border-radius: 6px; padding: 0 6px; }
.sd-hl { margin-left: auto; font-weight: 600; }
.sd-q { font-size: 0.9em; opacity: 0.8; margin: 4px 0 8px; }
.sd-row { display: grid; grid-template-columns: minmax(80px, 30%) 1fr 48px; gap: 8px; align-items: center; margin: 3px 0; }
.sd-name { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 0.92em; }
.sd-best { font-weight: 700; }
.sd-bar { background: var(--background-fill-secondary); border-radius: 4px; height: 12px; overflow: hidden; display: block; }
.sd-bar > span { display: block; height: 100%; background: var(--color-accent); border-radius: 4px; }
.sd-p { text-align: right; font-variant-numeric: tabular-nums; font-size: 0.9em; }
"""

with gr.Blocks(title="Super Decider") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Super Decider\n"
            "Send a **state** (text, JSON or images) and **typed questions**; get one calibrated answer per question: "
            "a `choice` with a probability per option, a `noul` (P(yes)), or a `score` over an ordered rubric. "
            "All answers are read in one denoising step of a fine-tuned DiffusionGemma 26B-A4B (NVFP4).\n\n"
            "[Model](https://huggingface.co/Singularity-Compute/super-decider) · "
            "[GitHub](https://github.com/Singularity-Compute/super-decider) · "
            "[API format](https://github.com/Singularity-Compute/super-decider/blob/main/docs/API.md)"
        )
        with gr.Row():
            with gr.Column(scale=1):
                state = gr.Textbox(label="State (text or JSON)", lines=6, value=EX_TICKET_STATE)
                questions = gr.Code(label="Questions (Jev / System One JSON)", language="json",
                                    value=_j(EX_TICKET_Q), lines=14)
                thinking = gr.Radio(["off", "auto", "always"], value="auto", label="Thinking",
                                    info="auto: think only when unsure (text, ≤2 questions); always: think first")
                images = gr.Gallery(label="Images (optional, up to 8)", type="filepath", interactive=True,
                                    columns=4, height=180)
                run = gr.Button("Decide", variant="primary")
                with gr.Accordion("Advanced", open=False):
                    seed = gr.Number(label="Seed (canvas noise)", value=42, precision=0)
            with gr.Column(scale=1):
                out_html = gr.HTML(label="Answers")
                out_json = gr.JSON(label="Response (Jev shape)")

        run.click(decide, inputs=[state, questions, thinking, images, seed], outputs=[out_html, out_json],
                  api_name="decide")

        gr.Examples(
            label="Text examples",
            examples=[
                [EX_TICKET_STATE, _j(EX_TICKET_Q), "auto"],
                [EX_INVOICE_STATE, _j(EX_INVOICE_Q), "always"],
                [_j(EX_ROUTE_STATE), _j(EX_ROUTE_Q), "auto"],
            ],
            inputs=[state, questions, thinking],
            outputs=[out_html, out_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )
        gr.Examples(
            label="Image example",
            examples=[[EX_IMAGE_STATE, _j(EX_IMAGE_Q), "off", ["examples/invoice-demo.png"]]],
            inputs=[state, questions, thinking, images],
            outputs=[out_html, out_json],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)