"""JEV-9B — typed decisions in one forward pass (System 1) + System 2 escalation.

Interactive demo for https://huggingface.co/autotrust/JEV-9B — AutoTrust's integrated
System 1 + System 2 model built with the Blocks-of-Experts recipe on a frozen
Qwen3.5-9B backbone. System 1 answers typed questions (`noul` yes/no, `choice`
over 2-16 options, `score` on a 0-5 scale) in one prefill pass with a calibrated
probability distribution; System 2 is the pristine base model (adapter off) and can
be escalated to for deliberate step-by-step reasoning.
"""

import os

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

import spaces  # noqa: E402  -- MUST precede torch / transformers / any CUDA-touching import

import html  # noqa: E402
import json  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402  # noqa: F401  (imported for dtype selection downstream)

import jev_core as jc  # noqa: E402

N_QUESTIONS = 4
LETTERS = jc.LETTERS
KINDS = list(jc.KINDS)
SCORE_LEVELS = [str(i) for i in range(6)]

TOPIC_OPTIONS = [
    "open-source software", "programming languages & tools", "AI / machine learning",
    "security & privacy", "hardware & retro computing", "science & health",
    "automotive & energy", "politics & government", "business & finance", "culture & lifestyle",
]
INTENT_OPTIONS = [
    "asking a question", "sharing personal experience", "providing factual information",
    "arguing an opinion", "attacking another commenter", "moderating the discussion",
]

print(
    f"[jev-9b] loaded {jc.REPO}: backbone dtype={next(jc._BACKBONE.parameters()).dtype}, "
    f"head dtype={jc._HEAD.weight.dtype}, temperatures={jc.TEMPERATURES}",
    flush=True,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _parse_options(text: str) -> list:
    """Options are comma-separated; newlines are accepted as separators too."""
    if not text:
        return []
    return [o.strip() for o in text.replace("\n", ",").split(",") if o.strip()]


def _build_questions(fields) -> list:
    """Turn the (type, instructions, options) UI rows into jev_core.decide calls."""
    questions = []
    for i in range(N_QUESTIONS):
        kind = (fields[3 * i] or "").strip()
        instructions = (fields[3 * i + 1] or "").strip()
        raw_options = fields[3 * i + 2] or ""
        if not instructions:
            continue  # a row with no instructions is skipped
        if kind not in KINDS:
            raise gr.Error(f"Question {i + 1}: unknown question type {kind!r}.")
        if kind == "choice":
            options = _parse_options(raw_options)
            if not 2 <= len(options) <= 16:
                raise gr.Error(
                    f"Question {i + 1} is a `choice` question, which needs 2-16 "
                    f"comma-separated options; got {len(options)}."
                )
        elif kind == "score":
            options = SCORE_LEVELS  # fixed 0-5 levels; the instructions describe the scale
        else:  # noul
            options = ["false", "true"]
        questions.append({"kind": kind, "instructions": instructions, "options": options})
    if not questions:
        raise gr.Error("Add at least one question — give it some instructions text.")
    return questions


def _bars_html(items: list, winner) -> str:
    """Render labelled probability bars; the winner is highlighted."""
    out = []
    for label, prob in items:
        hot = label == winner
        out.append(
            '<div style="display:flex;align-items:center;gap:10px;margin:2px 0">'
            f'<span style="flex:0 0 230px;text-align:right;font-size:12.5px;'
            f'color:var(--body-text-color);opacity:{1.0 if hot else 0.6};'
            f'overflow:hidden;text-overflow:ellipsis;white-space:nowrap">'
            f"{html.escape(str(label))}</span>"
            '<div style="flex:1;background:var(--border-color-primary);height:13px;border-radius:3px">'
            f'<div style="width:{max(prob * 100, 0.6):.2f}%;height:13px;border-radius:3px;'
            f'background:var(--color-accent);opacity:{1.0 if hot else 0.45}"></div></div>'
            f'<span style="flex:0 0 58px;font-size:12.5px;opacity:{1.0 if hot else 0.6}">'
            f"{prob * 100:.1f}%</span></div>"
        )
    return "".join(out)


def _answer_html(q: dict, probs: dict) -> str:
    """One question -> headline answer plus a labelled probability bar per option."""
    kind, head_txt = q["kind"], html.escape(q["instructions"])
    items = list(probs.items())
    if kind == "choice":
        winner = max(probs, key=probs.get)
        headline = html.escape(str(winner))
        note = f"confidence {probs[winner]:.3f}"
    elif kind == "score":
        ev = sum(int(k) * v for k, v in probs.items())
        winner = max(probs, key=probs.get)
        headline = f"expected level {ev:.2f} / 5"
        note = f"most likely level {winner} ({probs[winner] * 100:.1f}%)"
        items = [(f"level {k}", v) for k, v in items]
    else:  # noul
        p_yes = probs["true"]
        winner = "true" if p_yes >= 0.5 else "false"
        headline = "true" if p_yes >= 0.5 else "false"
        note = f"P(true) = {p_yes:.3f}"
    return (
        '<div style="padding:10px 0;border-bottom:1px solid var(--border-color-primary)">'
        f'<div style="font-size:12.5px;opacity:0.75;margin-bottom:3px">'
        f'<code style="font-size:11.5px">{kind}</code> {head_txt}</div>'
        f'<div style="font-size:19px;font-weight:600;margin:2px 0">{headline} '
        f'<span style="font-weight:400;font-size:13px;opacity:0.75">{note}</span></div>'
        f"{_bars_html(items, winner)}</div>"
    )


def _split_reasoning(text: str) -> tuple:
    """Split a System 2 reply into (reasoning, answer) if it carries think tags."""
    open_tag, close_tag = "<think>", "</think>"
    if close_tag in text:
        reasoning = text.split(close_tag)[0].replace(open_tag, "").strip()
        answer = text.split(close_tag, 1)[1].strip()
        return reasoning, answer
    return "", text.strip()


def _matched_option(answer: str, options: list):
    """Return the option the System 2 answer's last line picked, if any."""
    if not answer.strip():
        return None
    last = answer.strip().splitlines()[-1].strip().rstrip(".").strip("`").lower()
    for i, o in enumerate(options):
        ol = str(o).lower()
        if last == ol or last == f"{LETTERS[i]}) {ol}" or last == LETTERS[i].lower():
            return o
    for o in options:  # fall back to a contained mention
        if str(o).lower() in last:
            return o
    return None


# ---------------------------------------------------------------------------
# GPU handlers
# ---------------------------------------------------------------------------


@spaces.GPU(duration=45)
def decide(state: str, *fields):
    """Answer several typed questions about one piece of state text in one forward pass each.

    Args:
        state: the context to reason about — a customer message, a forum post, a diff.
        fields: four (type, instructions, options) triples. `type` is choice / score / noul;
            `options` is a comma-separated list (used by `choice`, ignored by the others).

    Returns:
        (answers as HTML, a one-line status, the raw answers as JSON)
    """
    state = (state or "").strip()
    if not state:
        raise gr.Error("Give the model some state text to reason about.")
    questions = _build_questions([f if isinstance(f, str) else "" for f in fields])

    state_used, truncated, n_tokens = jc.truncate_state(state)
    t0 = time.perf_counter()
    answers = [jc.decide(q["kind"], state_used, q["instructions"], q["options"])
               for q in questions]
    ms = (time.perf_counter() - t0) * 1000.0

    note = (f" — state truncated from {n_tokens} to {jc.MAX_STATE_TOKENS} tokens (head 60% / tail 40%)"
            if truncated else "")
    status = (
        f"**{len(questions)} question{'s' if len(questions) != 1 else ''} answered** — "
        f"{ms:.0f} ms of GPU compute for {len(questions)} prefill pass"
        f"{'es' if len(questions) != 1 else ''}, state {n_tokens} tokens{note}. "
        "System 1: one forward pass per question, no tokens generated."
    )
    payload = json.dumps(
        {"state": state, "system": 1, "answers": answers}, indent=2, ensure_ascii=False
    )
    return "".join(_answer_html(q, a) for q, a in zip(questions, answers)), status, payload


@spaces.GPU(duration=150)
def escalate(state: str, question: str, options_text: str, threshold: float,
             always_system2: bool, thinking: bool, max_new_tokens: float):
    """Confidence-gated escalation: System 1 answers in one pass; System 2 reasons when unsure.

    Args:
        state: the context to reason about.
        question: the question to answer.
        options_text: comma-separated options (2-16). Falls back to yes/no if blank.
        threshold: escalate to System 2 when System 1's top probability is below this.
        always_system2: skip the gate and always ask System 2 to reason.
        thinking: run System 2 in thinking mode (step-by-step reasoning).
        max_new_tokens: token budget for the System 2 reply.

    Returns:
        (answer HTML, status, System 2 reply, reasoning, raw JSON)
    """
    state = (state or "").strip()
    question = (question or "").strip()
    if not state:
        raise gr.Error("Give the model some state text to reason about.")
    if not question:
        raise gr.Error("Give the model a question.")
    options = _parse_options(options_text)
    kind = "choice" if options else "noul"
    if kind == "choice" and not 2 <= len(options) <= 16:
        raise gr.Error(f"`choice` needs 2-16 comma-separated options; got {len(options)}.")

    state_used, _trunc, n_tokens = jc.truncate_state(state)
    t0 = time.perf_counter()
    probs = jc.decide(kind, state_used, question, options if kind == "choice" else None)
    s1_ms = (time.perf_counter() - t0) * 1000.0
    top = max(probs, key=probs.get)
    gate = always_system2 or probs[top] < float(threshold)

    s1_html = (
        '<div style="padding:10px 0;border-bottom:1px solid var(--border-color-primary)">'
        f'<div style="font-size:12.5px;opacity:0.75;margin-bottom:3px">'
        f'<code style="font-size:11.5px">system 1</code> {html.escape(question)}</div>'
        f'<div style="font-size:19px;font-weight:600;margin:2px 0">{html.escape(str(top))} '
        f'<span style="font-weight:400;font-size:13px;opacity:0.75">'
        f"P({html.escape(str(top))}) = {probs[top]:.3f} — one prefill pass</span></div>"
        f"{_bars_html(list(probs.items()), top)}</div>"
    )

    if not gate:
        status = (
            f"**System 1 was confident** — top probability {probs[top]:.3f} ≥ threshold "
            f"{float(threshold):.2f}, so System 2 was not consulted. "
            f"({s1_ms:.0f} ms of GPU compute, state {n_tokens} tokens.)"
        )
        payload = json.dumps({"state": state, "question": question, "options": list(probs),
                              "system1": probs, "escalated": False}, indent=2, ensure_ascii=False)
        return s1_html, status, "", "", payload

    opts = list(probs.keys())
    t1 = time.perf_counter()
    raw, completed = jc.generate(state_used, question, opts,
                                 max_new_tokens=int(max_new_tokens),
                                 enable_thinking=bool(thinking))
    s2_ms = (time.perf_counter() - t1) * 1000.0
    reasoning, answer = _split_reasoning(raw)

    cut_note = ""
    if not completed:
        cut_note = (
            f"\n\n*System 2 hit its {int(max_new_tokens)}-token budget mid-reply — "
            "raise the budget, or untick thinking mode for a concise answer.*"
        )
        if thinking and not reasoning:
            # the verbose thinking trace never reached its closing tag
            reasoning, answer = raw, ""
    picked = _matched_option(answer, opts) if completed else None

    trigger = (
        f"the 'always ask System 2' box was ticked (System 1's top probability was "
        f"{probs[top]:.3f})"
        if always_system2 else
        f"System 1's top probability {probs[top]:.3f} was below the threshold "
        f"{float(threshold):.2f}"
    )
    status = (
        f"**Escalated to System 2** — {trigger}. System 1: "
        f"{s1_ms:.0f} ms; System 2: {s2_ms:.0f} ms, state {n_tokens} tokens."
        + (" Reply cut short at the token budget." if not completed else "")
    )
    answer_md = answer
    if picked:
        answer_md = (
            f"> **System 2 picks: `{picked}`**\n\n" + answer
        )
    answer_md += cut_note
    payload = json.dumps({"state": state, "question": question, "options": opts,
                          "system1": probs, "escalated": True, "completed": completed,
                          "system2_answer": answer, "system2_reasoning": reasoning},
                         indent=2, ensure_ascii=False)
    return s1_html, status, answer_md, reasoning, payload


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
main, .gradio-container, .gradio-container > .main { max-width: 1150px !important; margin: 0 auto !important; }
.dark .gradio-container { color: var(--body-text-color); }
"""

# Example rows, taken from the model authors' own published tests: the quickstart cases in
# the model card and the front-page / real-world cases in reports/realworld_9b.json
# (Apache-2.0). Each row is [state, (type, instructions, options) x 4].
EXAMPLES = [
    [
        "Customer says the parcel arrived damaged and wants their money back.",
        "noul", "Is the customer asking for a refund?", "",
        "noul", "Is the customer angry or frustrated?", "",
        "choice", "Which team should handle this message?",
        "billing, logistics & shipping, fraud & abuse, product feedback",
        "score", "How polite is the customer, from 0 (abusive) to 5 (very polite)?", "",
    ],
    [
        "SKU AX-330 stock at 8% of safety level; supplier late twice this quarter.",
        "choice", "Supplier response for this scenario.",
        "issue_warning, renegotiate, dual_source, maintain",
        "noul", "Is the supply risk already material?", "",
        "score", "Rate the supply risk on a 0-5 scale.", "",
        "noul", "", "",
    ],
    [
        "Two-tier encryption in the UK",
        "choice", "Which topic is this story mainly about?", ", ".join(TOPIC_OPTIONS),
        "noul", "Is this story about artificial intelligence or machine learning?", "",
        "noul", "", "",
        "noul", "", "",
    ],
    [
        "Please stop this. We've asked you before to observe the guidelines. HN is not a "
        "site for political debates, especially ideological ones. We've banned accounts "
        "repeatedly for this.",
        "choice", "What is this comment mainly doing?", ", ".join(INTENT_OPTIONS),
        "noul", "Does this comment insult or personally attack another person or group, "
                "rather than just arguing the point?", "",
        "noul", "", "",
        "noul", "", "",
    ],
    [
        "STATUS: payments processor is down. Checkout API returns 500s for all "
        "transactions; customers cannot pay. Revenue impact growing every minute.",
        "score", "Rate operational urgency on a 0-5 scale.", "",
        "noul", "Should the on-call engineer be paged now?", "",
        "noul", "", "",
        "noul", "", "",
    ],
]

with gr.Blocks(title="JEV-9B — typed decisions") as demo:
    gr.Markdown(
        "# JEV-9B — typed decisions in one forward pass\n"
        "An integrated **System 1 + System 2** model built with the Blocks-of-Experts recipe: "
        "a frozen **Qwen3.5-9B** backbone, a 40M-parameter LoRA + decision head for "
        "**System 1** — typed questions answered in a *single prefill pass* with a **calibrated "
        "probability distribution** — and the untouched base model as **System 2** for "
        "deliberate step-by-step reasoning. Per the authors' measurements: mean KL ≈ 0.019 to "
        "the closed TypeSafe Jev 1.13 on 25k held-out decisions, ECE 0.0007, and HumanEval "
        "70.7% — byte-identical to the base model.\n\n"
        "Model: [`autotrust/JEV-9B`](https://huggingface.co/autotrust/JEV-9B) (Apache-2.0) · "
        "System 2 path = [`Qwen/Qwen3.5-9B`](https://huggingface.co/Qwen/Qwen3.5-9B) unchanged"
    )

    with gr.Tabs():
        with gr.TabItem("System 1 — decide"):
            gr.Markdown(
                "Give the model one piece of **state** text and up to four typed questions. "
                "Nothing is generated: the model ranks the options you hand it and returns a "
                "probability distribution for each question — one forward pass per question."
            )
            state = gr.Textbox(
                label="State — the text to reason about",
                placeholder="A customer support message, a forum post, an incident summary… "
                            "up to 1,024 tokens (truncated head 60% / tail 40% beyond that)",
                lines=5, max_lines=14,
            )
            gr.Markdown(
                "#### Questions\n"
                "`choice` = one of 2–16 options you list · `score` = a fixed 0–5 ordered "
                "scale (describe the scale in the question) · `noul` = a yes/no statement, "
                "answered as *P(true)*. Leave a row's instructions blank to skip it."
            )
            kind_boxes, instr_boxes, opt_boxes = [], [], []
            for i in range(N_QUESTIONS):
                with gr.Row():
                    k = gr.Dropdown(choices=KINDS, value="noul", label=f"Q{i + 1} type",
                                    scale=1, allow_custom_value=True,
                                    info="leave the instructions blank to skip a row")
                    ins = gr.Textbox(label=f"Q{i + 1} instructions", scale=4, lines=1,
                                     placeholder="e.g. Which team should handle this?")
                    opt = gr.Textbox(label="options (comma-separated, choice only)", scale=5,
                                     lines=1, placeholder="e.g. billing, technical, sales")
                kind_boxes.append(k); instr_boxes.append(ins); opt_boxes.append(opt)

            run = gr.Button("Decide", variant="primary")
            out_status = gr.Markdown()
            out_html = gr.HTML(label="Answers")
            with gr.Accordion("Raw answers (JSON)", open=False):
                out_json = gr.Code(language="json", lines=14, show_label=False, wrap_lines=True)

            field_boxes = [c for triple in zip(kind_boxes, instr_boxes, opt_boxes) for c in triple]
            decide_inputs = [state] + field_boxes
            run.click(decide, inputs=decide_inputs,
                      outputs=[out_html, out_status, out_json], api_name="decide")

            gr.Examples(
                examples=EXAMPLES,
                inputs=decide_inputs,
                outputs=[out_html, out_status, out_json],
                fn=decide,
                cache_examples=True,
                cache_mode="lazy",
                label="Examples — from the model authors' published tests",
            )

        with gr.TabItem("System 1 → System 2 — escalate"):
            gr.Markdown(
                "The confidence-gated escalation pattern from the model card: System 1 answers "
                "in one pass; if its top probability is below your threshold, the *same "
                "weights* switch to System 2 — the pristine Qwen3.5-9B with the LoRA off — "
                "which thinks it through and writes an answer. Expect the two to disagree "
                "sometimes: they are integrated in serving, not in knowledge."
            )
            esc_state = gr.Textbox(
                label="State — the text to reason about", lines=5, max_lines=14,
                placeholder="The context your question is about…",
            )
            esc_question = gr.Textbox(
                label="Question", lines=1,
                placeholder="e.g. Which handler should receive this message?",
            )
            esc_options = gr.Textbox(
                label="Options (comma-separated, 2–16; leave blank for a yes/no question)",
                lines=1, placeholder="e.g. BLOCK, ESCALATE, ROUTE_TECH, ROUTE_SPAM",
            )
            with gr.Accordion("Escalation settings", open=True):
                esc_threshold = gr.Slider(
                    0.5, 1.0, value=0.90, step=0.01, label="Confidence threshold",
                    info="Escalate to System 2 when System 1's top probability is below this. "
                         "0.90 is the value suggested in the model card.",
                )
                esc_always = gr.Checkbox(
                    value=False, label="Always ask System 2 (skip the gate)",
                )
                esc_thinking = gr.Checkbox(
                    value=True, label="System 2 thinking mode (step-by-step reasoning)",
                )
                esc_tokens = gr.Slider(
                    64, 2048, value=1024, step=64, label="System 2 token budget",
                    info="Thinking mode is verbose — if the reply gets cut short, raise this.",
                )
            esc_btn = gr.Button("Ask", variant="primary")
            esc_status = gr.Markdown()
            esc_html = gr.HTML(label="System 1 — one prefill pass")
            esc_answer = gr.Markdown(label="System 2 — deliberate answer")
            with gr.Accordion("System 2 reasoning trace", open=False):
                esc_reasoning = gr.Markdown(label="Reasoning")
            with gr.Accordion("Raw (JSON)", open=False):
                esc_json = gr.Code(language="json", lines=12, show_label=False, wrap_lines=True)

            esc_btn.click(
                escalate,
                inputs=[esc_state, esc_question, esc_options, esc_threshold, esc_always,
                        esc_thinking, esc_tokens],
                outputs=[esc_html, esc_status, esc_answer, esc_reasoning, esc_json],
                api_name="escalate",
            )

    gr.Markdown(
        "---\n"
        "**What to expect.** System 1 is a distillation of the closed TypeSafe Jev 1.13's "
        "published output distributions: it mirrors the teacher's decisions, *including its "
        "mistakes*, and inherits its blind spots on arithmetic, dates, counting and "
        "adversarial inputs. It reads English best; the first 1,024 tokens of state are "
        "used. `choice` handles up to 16 options — beyond ~8, the authors recommend "
        "[JEV-27B](https://huggingface.co/autotrust/JEV-27B). Not for high-stakes decisions; "
        "use confidence gating.\n\n"
        "**Implementation.** Inference follows the authors' documented `transformers` + "
        "`peft` path exactly (backbone with the LoRA adapter on, 24-slot fp32 head on the "
        "last token, per-kind calibration temperature; System 2 generated with the adapter "
        "off). Weights: 18 GB bf16, streamed into VRAM by ZeroGPU on each request. "
        "Examples are taken from the model card quickstart and the authors' "
        "`reports/realworld_9b.json` (Apache-2.0)."
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)