"""decider-0.8b — typed decisions with calibrated probabilities, in one forward pass.

The model never generates text. It reads a *state* (a string or any JSON value) plus one or more
*typed questions* (Choice / Noul / Score, each with its own option list) and returns a calibrated
probability distribution per question, read off the hidden state at one answer slot per question.

The scoring path below is the authors' own reference implementation, vendored verbatim in
`decider/` (the inference subset shipped inside https://huggingface.co/Mapika/decider-0.8b,
Apache-2.0: `prompt.py`, `model.py`, `systemone.py`, `infer.py`). `_system_one` is
`decider.infer.Decider.system_one`'s eager branch with two changes only:

  * the checkpoint is loaded once at module scope and moved to CUDA eagerly (ZeroGPU), and
  * the temperature is an explicit argument instead of instance state, so concurrent requests
    on one worker cannot race each other.

`decider.engine` (shape-bucketed CUDA graphs) is deliberately not used: graph capture inside a
forked ZeroGPU worker buys little for a single interactive request. Everything else — the prompt
layout, the per-question row isolation, the isolated Score levels, the restriction of the answer
logits to the option-label tokens and the `-inf` masking of unused slots — is unmodified.
"""

import os

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

import spaces  # noqa: E402  — must precede torch / transformers

import json  # noqa: E402
import re  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402

from decider.infer import Example, Q, neutralize_options  # noqa: E402
from decider.model import DecisionModel, collate  # noqa: E402
from decider.prompt import MAX_OPTIONS, build  # noqa: E402
from decider.systemone import (  # noqa: E402
    MAX_CHOICE,
    MAX_LEVELS,
    assemble,
    plan_rows,
    render_question,
    render_state,
    unique_tokens,
)

MODEL_ID = "Mapika/decider-0.8b"

CFG = json.load(open(hf_hub_download(MODEL_ID, "decider_config.json")))
VERSION = CFG.get("version", "dev")
MODEL_NAME = "decider-" + str(VERSION)
DEFAULT_T = float(CFG.get("temperature", 1.0))           # 1.03 for 0.8b-v1 — the fitted, calibrated value
NEUTRALIZE_NONE = bool(CFG.get("neutralize_none", True))  # False for 0.8b-v1
DEFAULT_ISOLATED = bool(CFG.get("isolated_levels", False))
MAX_STATE_TOKENS = int(CFG.get("max_state_tokens", 32768))
MAX_FWD_TOKENS = 65536

# --------------------------------------------------------------------------------------------
# Model — loaded once at module scope, eagerly moved to the (ZeroGPU-intercepted) CUDA device.
# --------------------------------------------------------------------------------------------
MODEL = DecisionModel(MODEL_ID, dtype=torch.bfloat16, grad_ckpt=False).to("cuda").eval()
TOK = MODEL.tok


class _Keep:
    """The reference's rng stand-in: keep the caller's option order, never shuffle."""

    def shuffle(self, x):
        pass

    def sample(self, xs, k):
        return xs[:k]


def _system_one(state, questions, temperature, independent, isolated):
    """decider.infer.Decider.system_one, eager branch, with an explicit temperature."""
    ctx = render_state(state)
    rqs = {k: render_question(v) for k, v in questions.items()}
    opts = (lambda r: neutralize_options(r["options"])[0]) if NEUTRALIZE_NONE else (lambda r: list(r["options"]))
    flat, index = plan_rows(rqs, isolated and independent)
    rows = [[r] for r in flat] if independent else [flat]
    items = [
        build(
            Example(ctx, [Q(r["question"], opts(r), 0) for r in row]),
            TOK,
            _Keep(),
            max_options=MAX_OPTIONS,
            max_ctx_tokens=MAX_STATE_TOKENS,
            layout="state_first",
        )
        for row in rows
    ]
    probs = []
    with torch.no_grad():
        per = max(1, MAX_FWD_TOKENS // max(len(it["ids"]) for it in items))
        for i in range(0, len(items), per):
            chunk = items[i : i + per]
            bt = collate(chunk, TOK.pad_token_id)
            lg = MODEL.slot_logits(
                *[bt[k].to("cuda") for k in ("input_ids", "attention_mask", "slot_idx", "slot_batch", "nopts")]
            )
            pr = torch.softmax(lg / float(temperature), -1).float().cpu()
            c = 0
            for it in chunk:
                probs.append(pr[c : c + len(it["slots"])])
                c += len(it["slots"])
    flatp = [p.tolist() for ps in probs for p in ps]
    return (
        {
            "model": MODEL_NAME,
            "answers": assemble(rqs, index, flatp),
            "usage": {"input_tokens": unique_tokens(items), "output_tokens": 0},
        },
        items,
    )


# --------------------------------------------------------------------------------------------
# The question mini-format  ->  the Jev-shaped question dict the reference expects
# --------------------------------------------------------------------------------------------
TYPES = ("choice", "noul", "bool", "score")
_ID_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")


class ParseError(ValueError):
    pass


def _maybe_json(text):
    """Option descriptions may be a string or any JSON value (the reference renders both)."""
    t = text.strip()
    if t[:1] in "{[" or t in ("null", "true", "false"):
        try:
            return json.loads(t)
        except Exception:
            return text
    return text


def parse_questions(block: str) -> dict:
    """Parse the textarea into `{id: {"type", "instructions", "criteria"}}`.

    One question per blank-line-separated paragraph:

        <type> <id>: <instructions>
        - <option>            (choice: `name` or `name: description`; score: a level description)

    `<type>` is one of choice / noul / score and may be omitted (choice is assumed); `<id>` is
    optional and only ever used to key the answer — question ids are never shown to the model.
    """
    questions, auto = {}, 0
    for para in block.replace("\r\n", "\n").split("\n\n"):
        lines = [ln.rstrip() for ln in para.split("\n") if ln.strip() and not ln.lstrip().startswith("#")]
        if not lines:
            continue
        head_line, opt_lines = lines[0].strip(), []
        for ln in lines[1:]:
            s = ln.strip()
            if s.startswith(("-", "*", "•")):
                s = s[1:].strip()
            if s:
                opt_lines.append(s)

        qtype, qid, instructions = "choice", None, head_line
        if ":" in head_line:
            head, rest = head_line.split(":", 1)
            words = head.split()
            if 1 <= len(words) <= 2 and words[0].lower() in TYPES and rest.strip():
                qtype = words[0].lower()
                instructions = rest.strip()
                if len(words) == 2 and _ID_RE.match(words[1]):
                    qid = words[1]
        if not instructions:
            raise ParseError(f"Question without instructions: `{head_line}`")

        auto += 1
        qid = qid or f"q{auto}"
        if qid in questions:
            raise ParseError(f"Duplicate question id `{qid}`.")

        if qtype in ("noul", "bool"):
            crit = {}
            for s in opt_lines:
                if ":" in s:
                    k, v = s.split(":", 1)
                    if k.strip().lower() in ("true", "yes"):
                        crit["true"] = _maybe_json(v)
                    elif k.strip().lower() in ("false", "no"):
                        crit["false"] = _maybe_json(v)
            questions[qid] = {"type": "noul", "instructions": instructions, "criteria": crit or None}
        elif qtype == "score":
            if not 2 <= len(opt_lines) <= MAX_LEVELS:
                raise ParseError(
                    f"`{instructions}` is a score question with {len(opt_lines)} level(s); "
                    f"it needs an ordered list of 2 to {MAX_LEVELS} level descriptions."
                )
            questions[qid] = {"type": "score", "instructions": instructions, "criteria": list(opt_lines)}
        else:
            crit = {}
            for s in opt_lines:
                if ":" in s:
                    name, desc = s.split(":", 1)
                    name, desc = name.strip(), desc.strip()
                else:
                    name, desc = s, ""
                if not name:
                    continue
                crit[name] = _maybe_json(desc) if desc else None
            if not 2 <= len(crit) <= MAX_CHOICE:
                raise ParseError(
                    f"`{instructions}` is a choice question with {len(crit)} option(s); "
                    f"it needs 2 to {MAX_CHOICE}."
                )
            questions[qid] = {"type": "choice", "instructions": instructions, "criteria": crit}
    if not questions:
        raise ParseError("No question found. Write at least one question paragraph.")
    return questions


def parse_state(text: str):
    t = text.strip()
    if t[:1] in "{[":
        try:
            return json.loads(t)
        except Exception:
            pass
    return t


# --------------------------------------------------------------------------------------------
# Rendering
# --------------------------------------------------------------------------------------------
BAR_W = 26


def _bar(p: float) -> str:
    n = int(round(max(0.0, min(1.0, p)) * BAR_W))
    return "█" * n + "░" * (BAR_W - n)


def _render_answers(spec: dict, res: dict, elapsed: float) -> str:
    answers, usage = res["answers"], res["usage"]
    n = len(answers)
    out = [
        f"### {n} decision{'s' if n != 1 else ''} · `{res['model']}` · "
        f"{elapsed * 1000:.0f} ms · {usage['input_tokens']} input tokens · 0 tokens generated",
        "",
    ]
    for qid, a in answers.items():
        instructions = spec[qid]["instructions"]
        out.append(f"**`{qid}`** · *{instructions}*")
        if a["type"] == "noul":
            p = a["noul"]
            out.append(f"\n`noul` **P(yes) = {p:.3f}**\n")
            out.append(f"- `{_bar(p)}` {p * 100:5.1f}%  yes")
            out.append(f"- `{_bar(1 - p)}` {(1 - p) * 100:5.1f}%  no")
        elif a["type"] == "choice":
            out.append(
                f"\n`choice: {a['choice']}` — confidence **{a['confidence']:.3f}**, "
                f"certainty {a['certainty']:.3f}\n"
            )
            probs = sorted(a["probabilities"].items(), key=lambda kv: kv[1], reverse=True)
            for name, p in probs[:10]:
                mark = "  ←" if name == a["choice"] else ""
                out.append(f"- `{_bar(p)}` {p * 100:5.1f}%  {name}{mark}")
            if len(probs) > 10:
                rest = sum(p for _, p in probs[10:])
                out.append(f"- *{len(probs) - 10} further option(s), {rest * 100:.1f}% combined*")
        else:
            legend = a["legend"]
            top = max(a["probabilities"], key=lambda k: a["probabilities"][k])
            out.append(
                f"\n`score: {a['score']}` (expected level, 0–{len(legend) - 1}) — "
                f"most likely level **{top}**, p = {a['confidence']:.3f}, certainty {a['certainty']:.3f}\n"
            )
            for k, desc in legend.items():
                p = a["probabilities"][k]
                mark = "  ←" if k == top else ""
                out.append(f"- `{_bar(p)}` {p * 100:5.1f}%  **{k}** · {desc}{mark}")
            if "fit_mass" in a:
                fits = " ".join(f"`{k}`={v:.2f}" for k, v in a["level_fit"].items())
                out.append(
                    f"\n  *isolated levels — each judged alone: {fits} · "
                    f"fit mass {a['fit_mass']:.2f} (near 1 when exactly one level fits)*"
                )
        out.append("")
    return "\n".join(out)


# --------------------------------------------------------------------------------------------
# Inference handler
# --------------------------------------------------------------------------------------------
@spaces.GPU(duration=15)
def decide(
    state_text: str,
    questions_text: str,
    temperature: float = DEFAULT_T,
    independent: bool = True,
    isolated: bool = DEFAULT_ISOLATED,
) -> tuple:
    """Score typed questions against a state and return calibrated probabilities, in one pass.

    Args:
        state_text: the state the decisions are about — free text, or a JSON object / array.
        questions_text: one question per blank-line-separated paragraph. First line is
            `<type> <id>: <instructions>` with type in choice / noul / score (choice if omitted);
            each following `- ` line is one option (choice: `name` or `name: description`;
            score: one level description, lowest first).
        temperature: divides the option logits before the softmax. the value in decider_config.json (1.03) is the fitted,
            calibrated value shipped in the checkpoint's decider_config.json.
        independent: score every question in its own row, so adding, removing or reordering
            questions cannot change another answer. False packs them into one row (faster).
        isolated: judge every Score level in its own row, without its number or its neighbours,
            then normalise the per-level fits.

    Returns:
        A markdown report of the distributions, the raw /v1/systemone JSON response, and the
        exact prompt the model was shown.
    """
    try:
        spec = parse_questions(questions_text or "")
    except ParseError as e:
        return (f"**{e}**", "{}", "")
    except Exception as e:  # noqa: BLE001
        return (f"**Could not parse the questions.** `{type(e).__name__}: {e}`", "{}", "")

    state = parse_state(state_text or "")

    t0 = time.perf_counter()
    try:
        res, items = _system_one(state, spec, temperature, bool(independent), bool(isolated))
    except ValueError as e:
        return (f"**Invalid request.** `{e}`", "{}", "")
    elapsed = time.perf_counter() - t0

    prompt = ("\n\n" + "─" * 70 + "\n\n").join(TOK.decode(it["ids"]) for it in items)
    return (
        _render_answers(spec, res, elapsed),
        json.dumps(res, indent=2, ensure_ascii=False),
        prompt,
    )


# --------------------------------------------------------------------------------------------
# Examples — the authors' own, from the model card, github.com/Mapika/decider `examples/`
# and `decider/infer.py`'s demo block.
# --------------------------------------------------------------------------------------------
EX_TICKET_STATE = json.dumps(
    {
        "ticket": {
            "messages": [
                {"from": "customer", "text": "I was charged twice for order A-104. Please refund the duplicate."}
            ]
        },
        "refund_policy": "Duplicate charges are eligible for a refund.",
    },
    indent=2,
)
EX_TICKET_Q = """choice department: Which team should handle this?
- returns: Exchanges, refunds, wrong or damaged items
- billing: {"what": "Charges, invoices", "not_for": "delivery"}
- other

noul refund_requested: Does `ticket.messages[0].text` request a refund?

score frustration: How frustrated is the customer?
- calm
- frustrated
- very frustrated"""

EX_ROUTING_STATE = "third time I'm asking. the app logs me out every time. fix it today or I'm gone"
EX_ROUTING_Q = """choice action: What is the user trying to do?
- check_balance: View the account balance or recent transactions
- approve_transfer: {"what": "Approve or confirm a pending outgoing transfer", "not_for": "Asking about a transfer's status"}
- support: Get help with a problem, a complaint, or a question about the product

noul urgent: Does the message convey urgency?

score frustration: How frustrated is the user?
- calm
- irritated
- angry or threatening to leave"""

EX_BUG_STATE = (
    "The export button crashes the settings page in Safari. It works in Chrome, but a few of our "
    "customers only use Safari. Console says TypeError: undefined is not an object (exportBtn.js:41)."
)
EX_BUG_Q = """score severity: How severe is the reported issue?
- Cosmetic; no impact to functionality
- Broken or degraded feature, but a workaround exists
- Blocking issue; no workaround exists

score frustration: How frustrated is the customer?
- calm
- frustrated
- very frustrated

score actionable: How much does the report give an engineer to work with?
- no details
- some details (what happened)
- clear reproduction steps or error messages"""

EX_PLAIN_STATE = "My card was charged twice for the same purchase and I want the extra charge refunded."
EX_PLAIN_Q = """choice department: Which department should handle this?
- billing
- technical support
- sales

choice sentiment: What is the customer's sentiment?
- angry
- neutral
- happy

noul refund_action: Does this need a refund action?"""

EX_INCIDENT_STATE = (
    "Hi, since this morning our Stripe webhook integration stopped firing and our checkout is down for all "
    "customers. We are losing orders every minute and our partner launch is on Thursday. Also I think we got "
    "billed twice last week. If this is not fixed today we will have to look at other providers."
)
EX_INCIDENT_Q = """noul revenue_impacted: Revenue currently impacted?

choice business_impact: What business impact?
- none
- degraded
- outage

noul integration_issue: Integration issue present?

choice account_health: Account health status?
- healthy
- watch
- at risk

choice incident_scope: Which incident scope?
- single_account
- multi_account
- platform_wide

noul security_concern: Security concern present?

noul duplicate_charge: Duplicate charge reported?

score churn_likelihood: Churn likelihood level?
- none
- low
- medium
- high

noul human_attention: Human attention needed?

noul feature_request: Immediate feature request?"""

EX_GAME_STATE = """You are on a frozen lake grid (4x4). S = start, F = frozen (safe), H = hole (falling in ends the game), G = goal. The ice is not slippery here. Reach G.

Situation: Grid (your position in brackets):
[S] F  F  F 
 F  H  F  H 
 F  F  F  H 
 H  F  F  G 
You are at row 0, column 0."""
EX_GAME_Q = """choice action: What should you do right now?
- move left
- move down
- move right
- move up"""

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="decider-0.8b") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# decider-0.8b\n"
            "A language model that **does not generate text**. It reads a *state* and a set of *typed questions* "
            "— each with its own option list — and returns a **calibrated probability distribution per question**, "
            "read off the hidden state at one answer slot per question. One forward pass, no decoding, no output "
            "outside the options you defined.\n\n"
            f"`{VERSION}` · 0.75B params (Qwen3.5-0.8B-Base), bf16, on ZeroGPU — the smallest decider · "
            "[weights](https://huggingface.co/Mapika/decider-0.8b) · "
            "[GitHub](https://github.com/Mapika/decider) · "
            "[decider-2b demo](https://huggingface.co/spaces/hugging-apps/decider-2b-demo)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                state_box = gr.Textbox(
                    value=EX_TICKET_STATE,
                    lines=14,
                    max_lines=26,
                    label="State",
                    info="Free text, or a JSON object / array. Up to 32k tokens.",
                )
            with gr.Column(scale=1):
                questions_box = gr.Textbox(
                    value=EX_TICKET_Q,
                    lines=14,
                    max_lines=26,
                    label="Typed questions",
                    info="One per blank-line-separated paragraph. First line: "
                    "'<type> <id>: <instructions>' with type = choice | noul | score. "
                    "Each '- ' line is one option ('name: description' for choice, "
                    "one level description for score, lowest first).",
                )
        run = gr.Button("Decide", variant="primary")

        answers_md = gr.Markdown("Answers appear here.")

        with gr.Accordion("Raw /v1/systemone response", open=False):
            raw_json = gr.Code(value="{}", language="json", label="JSON")
        with gr.Accordion("The prompt the model was shown", open=False):
            prompt_view = gr.Code(
                value="",
                label="One row per question — the model reads the logits at each 'Answer: (' slot",
            )

        with gr.Accordion("Advanced settings", open=False):
            temperature = gr.Slider(
                0.5,
                3.0,
                value=DEFAULT_T,
                step=0.01,
                label="Temperature",
                info=f"Divides the option logits before the softmax. {DEFAULT_T:.2f} is the value fitted on "
                "in-task data and shipped in decider_config.json — the calibrated setting.",
            )
            independent = gr.Checkbox(
                value=True,
                label="Independent questions",
                info="Score each question in its own row, so adding, removing or reordering questions cannot "
                "change another answer. Unchecking packs them into one row (about half the latency, but later "
                "questions then see the earlier question texts).",
            )
            isolated = gr.Checkbox(
                value=DEFAULT_ISOLATED,
                label="Isolated score levels",
                info="Judge every Score level in its own yes/no row, without its number or its neighbours, then "
                "normalise the per-level fits. This is the v8+ default.",
            )

        gr.Examples(
            examples=[
                [EX_TICKET_STATE, EX_TICKET_Q],
                [EX_ROUTING_STATE, EX_ROUTING_Q],
                [EX_BUG_STATE, EX_BUG_Q],
                [EX_PLAIN_STATE, EX_PLAIN_Q],
                [EX_INCIDENT_STATE, EX_INCIDENT_Q],
                [EX_GAME_STATE, EX_GAME_Q],
            ],
            inputs=[state_box, questions_box],
            outputs=[answers_md, raw_json, prompt_view],
            fn=decide,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples — from the model card, the repo's examples/ programs and its game envs",
        )

        gr.Markdown(
            "### Notes\n"
            "- **Calibrated, not chatty.** The top probability is meant to be a real confidence: the model was "
            "trained with a proper scoring rule and a single temperature fitted on in-task data (reported ECE "
            "0.03 in-task, 0.10 held-out). Ask it something open-ended and you still only get a distribution over *your* options.\n"
            "- **A 0.8B model without reasoning.** It gives up knowledge rather than the decision format: "
            "routing, yes/no and short-state lookups are within 1–4 points of the 2B, while knowledge-heavy multiple choice is close to the base "
            "model, and a judgment that needs several steps should be split into several questions.\n"
            "- **Rules written into a question** (\"fill if empty, otherwise skip\") are not followed at this "
            "size — state the decision as a plain question with described options.\n"
            "- Choice takes 2–255 options, Score 2–10 ordered levels, Noul returns P(yes). English only."
        )

        run.click(
            fn=decide,
            inputs=[state_box, questions_box, temperature, independent, isolated],
            outputs=[answers_md, raw_json, prompt_view],
            api_name="decide",
        )

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