"""Standalone in-weights runner for `Qwen3.5-9B-Math-GraphCalc` (and `0.8B`).

Requires `qwenputer` (`https://github.com/divagr18/qwenputer`) on `PYTHONPATH` or pointed to by
`QWENPUTER_ROOT` (defaults to `./qwenputer` or `../qwenputer`):

    git clone https://github.com/divagr18/qwenputer
    QWENPUTER_ROOT=./qwenputer python runner.py --ask "Calculate 37429 * 8613 + 194820 and explain your steps."
"""

from __future__ import annotations

import argparse
import copy
import gc
import json
import os
import re
import sys
import time
import zlib
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors import safe_open
from tokenizers import Tokenizer
from transformers import AutoConfig, AutoTokenizer, DynamicCache, Qwen3_5ForCausalLM
from transformers.models.qwen3_5 import modeling_qwen3_5 as q35_mod

MODEL_DIR = Path(__file__).resolve().parent


def _locate_qwenputer_rv() -> Path:
    candidates = []
    if os.environ.get("QWENPUTER_ROOT"):
        candidates.append(Path(os.environ["QWENPUTER_ROOT"]))
    candidates.extend([
        Path.cwd() / "qwenputer",
        Path.cwd().parent / "qwenputer",
        MODEL_DIR / "qwenputer",
        MODEL_DIR.parent / "qwenputer",
        MODEL_DIR.parent.parent / "qwenputer",
        MODEL_DIR.parent.parent.parent / "qwenputer",
    ])
    for root in candidates:
        rv = root.resolve() / "experimental" / "riscv"
        if (rv / "chipformer").is_dir():
            return rv
    raise RuntimeError(
        "Could not locate `chipformer` from `qwenputer` (`experimental/riscv/chipformer`). "
        "Clone https://github.com/divagr18/qwenputer and set QWENPUTER_ROOT=/path/to/qwenputer."
    )


try:
    import chipformer  # noqa: F401
except ImportError:
    sys.path.insert(0, str(_locate_qwenputer_rv()))

from chipformer.q35 import fast as F_chip  # noqa: E402
from chipformer.q35 import ir  # noqa: E402
from chipformer.q35.ir import Target  # noqa: E402
from chipformer.rv import bridge as BR  # noqa: E402
from chipformer.rv import core, elf, hf, ref  # noqa: E402
from chipformer.rv import spec as S  # noqa: E402

CARBON_SPECIALS = [
    "<tool_response>",   # 256
    "</tool_response>",  # 257
    "<think>",           # 258
    "</think>",          # 259
    "<|im_end|>",        # 260
    "<|im_start|>",      # 261
    "user",              # 262
    "assistant",         # 263
    "\n\n",              # 264
]

GRAPHCALC_TOOL = {
    "type": "function",
    "function": {
        "name": "graphcalc",
        "description": (
            "Exact in-weights graphing & math calculator (62-digit integers/rationals, scientific "
            "functions, equation solver, and coordinate plotter). Pass an arithmetic expression, "
            "an equation in x (LHS = RHS), or a function f(x); optionally pass xmin and xmax to "
            "plot y = f(x) and compute roots, vertex/extrema, definite integral, and value table."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Math expression, equation (LHS = RHS), or function of x (e.g. 'x^2 - 4*x + 3')",
                },
                "xmin": {
                    "type": "number",
                    "description": "Optional domain/plot minimum x (default -5 for plots, -20 for nonlinear solve)",
                },
                "xmax": {
                    "type": "number",
                    "description": "Optional domain/plot maximum x (default 5 for plots, 20 for nonlinear solve)",
                },
            },
            "required": ["expression"],
        },
    },
}

GRAPHCALC_SYSTEM_PROMPT = (
    "You are a clear, encouraging step-by-step math tutor equipped with an embedded graphing "
    "calculator (`graphcalc`). Always call `graphcalc` first to compute exact arithmetic, roots, "
    "coordinate tables, or function plots, then explain every step clearly so a student can learn."
)

_orig_chunk_gdn = q35_mod.torch_chunk_gated_delta_rule
_orig_rec_gdn = q35_mod.torch_recurrent_gated_delta_rule


def _fp32_chunk_gdn(*args, **kwargs):
    with torch.autocast("cuda", enabled=False):
        return _orig_chunk_gdn(*args, **kwargs)


def _fp32_rec_gdn(*args, **kwargs):
    with torch.autocast("cuda", enabled=False):
        return _orig_rec_gdn(*args, **kwargs)


q35_mod.torch_chunk_gated_delta_rule = _fp32_chunk_gdn
q35_mod.torch_recurrent_gated_delta_rule = _fp32_rec_gdn
F_chip.exact_matmuls()


class CarbonText(BR.Text):
    """Qwen3.5 bridge vocabulary with native ChatML special tokens (256..264) and space HALT echo."""

    def __init__(self, tokenizer_path: str | Path) -> None:
        p = Path(tokenizer_path)
        tj = p if p.is_file() else (p / "tokenizer.json")
        self.tok = Tokenizer.from_file(str(tj))
        vocab = self.tok.get_vocab(with_added_tokens=True)
        base_vocab = {s: i for s, i in vocab.items() if not s.startswith("<|chip_")}
        self.n = max(base_vocab.values()) + 1
        dec = {c: b for b, c in BR.bytes_to_unicode().items()}
        added = (
            {a.content for a in self.tok.get_added_tokens_decoder().values()}
            if hasattr(self.tok, "get_added_tokens_decoder")
            else set()
        )
        self.bytes = [b""] * self.n
        for s, i in base_vocab.items():
            if s in added or not all(ch in dec for ch in s):
                self.bytes[i] = s.encode("utf-8")
            else:
                self.bytes[i] = bytes(dec[ch] for ch in s)
        self.TCO = base_vocab["<tool_call>"]
        self.TCC = base_vocab["</tool_call>"]
        u = BR.bytes_to_unicode()
        special_ids = []
        for sp in CARBON_SPECIALS:
            if sp in base_vocab:
                special_ids.append(base_vocab[sp])
            else:
                enc = self.tok.encode(sp, add_special_tokens=False).ids
                assert len(enc) == 1, f"Special {sp!r} is not a single Qwen token: {enc}"
                special_ids.append(enc[0])
        self.printable = [base_vocab[u[b]] for b in range(256)] + special_ids
        self.pindex = {q: p for p, q in enumerate(self.printable)}
        self.NL = self.printable[32]


def _compact_layer_inplace(layer: nn.Module, is_chip: bool) -> tuple[int, int]:
    before = after = 0
    if is_chip and hasattr(layer, "linear_attn") and isinstance(layer.linear_attn.out_proj, nn.Linear):
        F_chip.prune_heads(layer.linear_attn)
    for parent in list(layer.modules()):
        for name, child in list(parent.named_children()):
            if isinstance(child, nn.Linear) and child.bias is None:
                sp = F_chip.SparseLinear(child.weight)
                before += child.weight.numel()
                after += sp.weight.numel()
                setattr(parent, name, sp)
    return before, after


def _load_and_compact_qwen35(path: Path, n_chip: int, dev: str) -> tuple[Qwen3_5ForCausalLM, dict]:
    layout_file = path / "chip_layout.json"
    if not layout_file.is_file():
        model = Qwen3_5ForCausalLM.from_pretrained(
            str(path),
            dtype=torch.bfloat16,
            attn_implementation="sdpa",
        )
        chip_b = chip_a = qwen_b = qwen_a = 0
        for i, layer in enumerate(model.model.layers):
            b, a = _compact_layer_inplace(layer, is_chip=(i < n_chip))
            if i < n_chip:
                chip_b += b
                chip_a += a
            else:
                qwen_b += b
                qwen_a += a
        return model.to(dev).eval(), {
            "chip_dense": chip_b,
            "chip_compact": chip_a,
            "qwen_dense": qwen_b,
            "qwen_compact": qwen_a,
        }

    layout = json.loads(layout_file.read_text(encoding="utf-8"))
    blob = (path / "chip_indices.bin.zlib").read_bytes()
    cfg = AutoConfig.from_pretrained(str(path))
    cfg._attn_implementation = "sdpa"
    with torch.device("meta"):
        model = Qwen3_5ForCausalLM(cfg)

    tensors_spec: dict[str, dict] = layout["tensors"]
    by_layer: list[list[str]] = [[] for _ in range(cfg.num_hidden_layers)]
    for k in tensors_spec:
        if k.startswith("model.layers."):
            lid = int(k.split(".", 3)[2])
            by_layer[lid].append(k)

    chip_before = chip_after = qwen_before = qwen_after = head_before = head_after = 0
    with safe_open(str(path / "model.safetensors"), framework="pt", device="cpu") as f:
        def _unpack_one(k: str) -> torch.Tensor:
            spec = tensors_spec[k]
            vals = f.get_tensor(k)
            if spec["mode"] == "dense":
                return vals
            shape = tuple(spec["shape"])
            raw = zlib.decompress(blob[spec["offset"] : spec["offset"] + spec["length"]])
            idx = np.cumsum(np.frombuffer(raw, dtype=np.uint32), dtype=np.int64)
            out = torch.zeros(int(np.prod(shape)), dtype=vals.dtype)
            out[torch.from_numpy(idx)] = vals
            return out.view(shape)

        for lid in range(cfg.num_hidden_layers):
            prefix = f"model.layers.{lid}."
            sub_sd = {k[len(prefix) :]: _unpack_one(k) for k in by_layer[lid]}
            layer = model.model.layers[lid]
            layer.load_state_dict(sub_sd, assign=True, strict=True)
            del sub_sd
            b, a = _compact_layer_inplace(layer, is_chip=(lid < n_chip))
            if lid < n_chip:
                chip_before += b
                chip_after += a
            else:
                qwen_before += b
                qwen_after += a
            layer.to(dev)

        norm_w = _unpack_one("model.norm.weight").to(dev)
        model.model.norm.load_state_dict({"weight": norm_w}, assign=True, strict=True)

        emb_w = _unpack_one("model.embed_tokens.weight").to(dev)
        model.model.embed_tokens.load_state_dict({"weight": emb_w}, assign=True, strict=True)

        if "lm_head.weight" in tensors_spec and not getattr(cfg, "tie_word_embeddings", True):
            head_w_cpu = _unpack_one("lm_head.weight")
            head_before = head_w_cpu.numel()
            hc0, hc1 = F_chip._hull(F_chip.nonzero_cols(head_w_cpu))
            hblock = nn.Parameter(head_w_cpu[:, hc0:hc1].contiguous().to(dev), requires_grad=False)
            head_after = hblock.numel()
            model.lm_head = F_chip.SparseLinear(
                head_w_cpu, rows=(0, head_w_cpu.shape[0]), cols=(hc0, hc1), block=hblock
            )
            del head_w_cpu
        else:
            model.lm_head.load_state_dict({"weight": emb_w}, assign=True, strict=True)
            model.tie_weights()
            head_before = head_after = emb_w.numel()

    model.model.rotary_emb = q35_mod.Qwen3_5TextRotaryEmbedding(config=cfg).to(dev)
    model.eval()
    if dev == "cuda":
        with torch.no_grad():
            for _ in range(2):
                _wc = DynamicCache(config=cfg)
                model.model(
                    input_ids=torch.zeros((1, 2), dtype=torch.long, device=dev),
                    past_key_values=_wc,
                    use_cache=True,
                )
                model.model(
                    input_ids=torch.zeros((1, 1), dtype=torch.long, device=dev),
                    past_key_values=_wc,
                    use_cache=True,
                )
                del _wc
        gc.collect()
        torch.cuda.empty_cache()
    else:
        gc.collect()
    return model, {
        "chip_dense": chip_before,
        "chip_compact": chip_after,
        "qwen_dense": qwen_before,
        "qwen_compact": qwen_after,
        "head_dense": head_before,
        "head_compact": head_after,
    }


class CarbonGraphCalcRunner:
    """Self-contained runner for `Qwen3.5-Math-GraphCalc` checkpoints."""

    def __init__(
        self,
        model_dir: str | Path = MODEL_DIR,
        firmware_path: str | Path | None = None,
        device: str | None = None,
        fast: bool = True,
        executor: bool | str = True,
    ) -> None:
        torch.set_grad_enabled(False)
        F_chip.exact_matmuls()
        self.path = Path(model_dir)
        self.meta = json.loads((self.path / "chipformer.json").read_text(encoding="utf-8"))
        self.dev = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.n_chip = int(self.meta["chip_layers"])
        self.v0 = int(self.meta["chip_token_offset"])
        self.d0 = int(self.meta["chip_dims_offset"])

        self.model, self.compact_stats = _load_and_compact_qwen35(self.path, self.n_chip, self.dev)
        self.lm = self.model.model
        if isinstance(self.model.lm_head, F_chip.SparseLinear):
            self.head_c0 = self.model.lm_head.c0
            self.head_c1 = self.model.lm_head.c1
            self.qwen_head_w = self.model.lm_head.weight[: self.v0]
        else:
            self.head_c0 = 0
            self.head_c1 = self.model.config.hidden_size
            self.qwen_head_w = self.model.lm_head.weight[: self.v0]
        self._p4_buf = torch.zeros((4, 1, 1), dtype=torch.long, device=self.dev)

        self.tok = AutoTokenizer.from_pretrained(str(self.path))
        self.T = CarbonText(self.path / "tokenizer.json")

        cfg = self.model.config
        self.c = core.build(
            Target(cfg.linear_num_key_heads, cfg.linear_num_value_heads, cfg.intermediate_size),
            self.T,
        )
        self.code_dims = torch.tensor(self.meta["code_dims"], device=self.dev, dtype=torch.long)
        self.eos = {248046, 248044}
        self.go_model_id = self.v0 + S.tok("GO")

        fw = firmware_path or (self.path / self.meta.get("firmware", "graphcalc.elf"))
        self.firmware_path = Path(fw)
        self.prog = elf.program(self.firmware_path)
        core.snapshot(self.c, self.fresh_machine())

        self.executor = executor if self.dev == "cuda" else False
        self.verify = self.executor == "verify"
        if self.executor:
            try:
                self._build_exec()
            except Exception as exc:
                if self.verify:
                    raise
                print(f"[runner] Triton Executor2 unavailable ({type(exc).__name__}: {exc}); using CUDA-graphed HF chip layers.", flush=True)
                self.executor = False
                self.verify = False

        self.fast = bool(fast and self.dev == "cuda" and (not self.executor or self.verify))
        if self.fast:
            self._build_fast()

    def _build_exec(self) -> None:
        from chipformer.rv.chat import ExecRow

        self.keep_ir = torch.tensor(self.c.keep, device=self.dev, dtype=torch.long)
        cols = self.meta["code_dims"] + [self.meta["tau_dim"], self.meta["mu_dim"]]
        self.keep_cols = torch.tensor(cols, device=self.dev, dtype=torch.long)
        self.emb_cache: dict[int, torch.Tensor] = {}
        self.erow = ExecRow(self)
        self.erow._capture()

    def _exec_emb(self, ids: list[int]) -> torch.Tensor:
        new = sorted({i for i in ids if i not in self.emb_cache})
        if new:
            rows = self.erow.ex.embed(
                [i - self.v0 if i >= self.v0 else BR.text_token(i) for i in new]
            )
            for i, r in zip(new, rows):
                self.emb_cache[i] = r
        return torch.stack([self.emb_cache[i] for i in ids])

    @torch.no_grad()
    def _exec_step_one(self, tok_id: int) -> torch.Tensor:
        row = self.emb_cache.get(tok_id)
        if row is None:
            st = tok_id - self.v0 if tok_id >= self.v0 else BR.text_token(tok_id)
            row = self.erow.ex.embed([st])[0]
            self.emb_cache[tok_id] = row
        self.erow.xbuf[0].copy_(row)
        self.erow.graph.replay()
        return self.erow.kbuf

    @torch.no_grad()
    def _exec_start_collect(self, ids: list[int]) -> torch.Tensor:
        if self.erow.graph is None:
            self.erow._capture()
        for st in self.erow.heads:
            st["S"].zero_()
            st["hist"].zero_()
        emb = self._exec_emb(ids)
        keeps = torch.empty((len(ids), len(self.keep_cols)), dtype=torch.float32, device=self.dev)
        for t in range(len(ids)):
            self.erow.xbuf.copy_(emb[t : t + 1])
            self.erow.graph.replay()
            keeps[t].copy_(self.erow.kbuf)
        self.erow.ex.load_init()
        return keeps

    def _exec_hc(self, tok_id: int, keep: torch.Tensor) -> torch.Tensor:
        hc = torch.zeros(self.lm.embed_tokens.weight.shape[1], dtype=torch.bfloat16, device=self.dev)
        hc[: self.d0] = self.lm.embed_tokens.weight[tok_id, : self.d0]
        hc[self.keep_cols] = keep.to(torch.bfloat16)
        return hc

    def _build_fast(self) -> None:
        cfg = copy.deepcopy(self.model.config)
        cfg.num_hidden_layers = self.n_chip + 1
        cfg.layer_types = ["linear_attention"] * self.n_chip + ["full_attention"]
        self.ccfg = cfg
        dummy = q35_mod.Qwen3_5DecoderLayer(cfg, self.n_chip).to(self.dev, torch.bfloat16)
        for p in dummy.parameters():
            p.data.zero_()
        with torch.device("meta"):
            cm = q35_mod.Qwen3_5TextModel(cfg)
        cm.embed_tokens = self.lm.embed_tokens
        cm.layers = nn.ModuleList([*self.lm.layers[: self.n_chip], dummy])
        cm.norm = nn.Identity()
        cm.rotary_emb = q35_mod.Qwen3_5TextRotaryEmbedding(config=cfg).to(self.dev)
        self.cm = cm.eval()
        self.graph = None

    def _chip_cache(self):
        if getattr(self, "ccache", None) is None:
            self.ccache = DynamicCache(config=self.ccfg)
        else:
            for L in range(self.n_chip):
                lay = self.ccache.layers[L]
                lay.conv_states[0].zero_()
                lay.recurrent_states[0].zero_()
                lay.has_previous_state[0] = False
        return self.ccache

    @torch.no_grad()
    def _chip_body(self):
        h = self.cm.embed_tokens(self.g_ids)
        for L in range(self.n_chip):
            h = self.cm.layers[L](
                h,
                position_embeddings=None,
                attention_mask=None,
                position_ids=None,
                past_key_values=self.ccache,
            )
        self.g_out.copy_(h)

    @torch.no_grad()
    def _capture(self):
        self.g_ids = torch.zeros((1, 1), dtype=torch.long, device=self.dev)
        self.g_out = torch.zeros(
            (1, 1, self.model.config.hidden_size), dtype=torch.bfloat16, device=self.dev
        )
        lays = [self.ccache.layers[L] for L in range(self.n_chip)]
        saved = [(lay.conv_states[0].clone(), lay.recurrent_states[0].clone()) for lay in lays]
        s = torch.cuda.Stream()
        s.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(s):
            for _ in range(3):
                self._chip_body()
        torch.cuda.current_stream().wait_stream(s)
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph):
            self._chip_body()
        for lay, (cs, rs) in zip(lays, saved):
            lay.conv_states[0].copy_(cs)
            lay.recurrent_states[0].copy_(rs)

    @torch.no_grad()
    def _chip_step(self, tok_id: int, g: int, cache, compiled: bool = True) -> torch.Tensor:
        if not compiled or self.graph is None:
            x = torch.tensor([[tok_id]], device=self.dev)
            out = self.cm(
                input_ids=x,
                past_key_values=cache,
                use_cache=True,
                cache_position=torch.tensor([g], device=self.dev),
            )
            return out.last_hidden_state[0, -1].clone()
        self.g_ids.fill_(tok_id)
        self.graph.replay()
        return self.g_out[0, -1].clone()

    @torch.no_grad()
    def _qwen_step(self, hc: torch.Tensor, cache, pos: int, full_head: bool = False) -> torch.Tensor:
        h = hc.view(1, 1, -1)
        p4 = self._p4_buf
        p4.fill_(pos)
        pe = self.lm.rotary_emb(h, p4[1:])
        for i in range(self.n_chip, len(self.lm.layers)):
            h = self.lm.layers[i](
                h,
                position_embeddings=pe,
                attention_mask=None,
                position_ids=p4[0],
                past_key_values=cache,
            )
        nh = self.lm.norm(h)[0, -1]
        if full_head:
            return self.model.lm_head(nh)
        return F.linear(nh[self.head_c0 : self.head_c1], self.qwen_head_w)

    @torch.no_grad()
    def _qwen_prefill(
        self, x_ids: torch.Tensor, keeps: torch.Tensor, full_head: bool = False
    ) -> tuple[DynamicCache, torch.Tensor]:
        p_len = int(x_ids.shape[1])
        h = torch.zeros(
            (1, p_len, self.model.config.hidden_size), dtype=torch.bfloat16, device=self.dev
        )
        h[0, :, : self.d0] = self.lm.embed_tokens.weight[x_ids[0], : self.d0]
        h[0, :, self.keep_cols] = keeps.to(torch.bfloat16)
        cache = DynamicCache(config=self.model.config)
        pos_vec = torch.arange(0, p_len, device=self.dev, dtype=torch.long).unsqueeze(0)
        p4 = pos_vec.unsqueeze(0).expand(4, -1, -1)
        pe = self.lm.rotary_emb(h, p4[1:])
        for i in range(self.n_chip, len(self.lm.layers)):
            h = self.lm.layers[i](
                h,
                position_embeddings=pe,
                attention_mask=None,
                position_ids=p4[0],
                past_key_values=cache,
            )
        nh = self.lm.norm(h)[0, -1]
        logits = (
            self.model.lm_head(nh)
            if full_head
            else F.linear(nh[self.head_c0 : self.head_c1], self.qwen_head_w)
        )
        return cache, logits

    @torch.no_grad()
    def _qwen_chunk(
        self, hcs: list[torch.Tensor], cache, pos0: int, full_head: bool = False
    ) -> torch.Tensor:
        if len(hcs) == 1:
            return self._qwen_step(hcs[0], cache, pos0, full_head=full_head)
        k_len = len(hcs)
        h = torch.stack(hcs, dim=0).unsqueeze(0)
        pos_vec = torch.arange(pos0, pos0 + k_len, device=self.dev, dtype=torch.long).unsqueeze(0)
        p4 = pos_vec.unsqueeze(0).expand(4, -1, -1)
        pe = self.lm.rotary_emb(h, p4[1:])
        attn_mask = torch.zeros((1, 1, k_len, pos0 + k_len), dtype=h.dtype, device=self.dev)
        future = torch.triu(torch.ones((k_len, k_len), dtype=torch.bool, device=self.dev), diagonal=1)
        attn_mask[0, 0, :, pos0:][future] = torch.finfo(h.dtype).min
        for i in range(self.n_chip, len(self.lm.layers)):
            lmask = attn_mask if self.model.config.layer_types[i] == "full_attention" else None
            h = self.lm.layers[i](
                h,
                position_embeddings=pe,
                attention_mask=lmask,
                position_ids=p4[0],
                past_key_values=cache,
            )
        nh = self.lm.norm(h)[0, -1]
        if full_head:
            return self.model.lm_head(nh)
        return F.linear(nh[self.head_c0 : self.head_c1], self.qwen_head_w)

    @torch.no_grad()
    def _layers_slow(self, tok_id: int, cache, pos: int | None):
        h = self.lm.embed_tokens(torch.tensor([[tok_id]], device=self.dev))
        pe = tp = None
        if pos is not None:
            p4 = torch.full((4, 1, 1), pos, device=self.dev, dtype=torch.long)
            tp = p4[0]
            pe = self.lm.rotary_emb(h, p4[1:])
        hc = None
        n = self.n_chip if pos is None else len(self.lm.layers)
        for i in range(n):
            h = self.lm.layers[i](
                h,
                position_embeddings=pe,
                attention_mask=None,
                position_ids=tp,
                past_key_values=cache,
            )
            if i == self.n_chip - 1:
                hc = h
        logits = None if pos is None else self.model.lm_head(self.lm.norm(h))[0, -1]
        return hc[0, -1], logits

    def _decode_code_list(self, vals: list[float]) -> int | None:
        if not any(v != 0.0 for v in vals):
            return None
        nf = core.NFAM
        fid = 0
        for i in range(nf):
            if vals[i] > 0.0:
                fid |= 1 << i
        p = 0
        for i in range(len(vals) - nf):
            if vals[nf + i] > 0.0:
                p |= 1 << i
        t = core.decode_code(fid, p, self.T)
        if t < 0:
            raise RuntimeError(f"Invalid circuit token prediction (family={fid}, payload={p})")
        return t - S.TEXT_TOKEN if S.is_text(t) else self.v0 + t

    def code_token(self, code: torch.Tensor) -> int | None:
        return self._decode_code_list(code.tolist())

    def circuit_token(self, hc: torch.Tensor) -> int | None:
        return self.code_token(hc[self.code_dims])

    def fresh_machine(self) -> ref.Machine:
        m = ref.Machine.from_program(self.prog)
        BR.park(m)
        return m

    def prompt_ids(
        self,
        user_prompt: str | list[dict],
        system_prompt: str | None = GRAPHCALC_SYSTEM_PROMPT,
        use_tools: bool = True,
    ) -> list[int]:
        if isinstance(user_prompt, str):
            msgs = []
            if system_prompt:
                msgs.append({"role": "system", "content": system_prompt})
            msgs.append({"role": "user", "content": user_prompt})
        else:
            msgs = list(user_prompt)
        tools = [GRAPHCALC_TOOL] if use_tools else None
        out = self.tok.apply_chat_template(
            msgs,
            tools=tools,
            add_generation_prompt=True,
            enable_thinking=False,
            tokenize=True,
        )
        if not isinstance(out, list):
            out = out["input_ids"]
        return [int(x) for x in out]

    @torch.no_grad()
    def generate(
        self,
        ids: list[int],
        max_new_text: int = 512,
        max_chip_tokens: int = 500_000,
        max_calls: int = 1,
        check_head: bool = False,
        check_lock: bool = False,
    ) -> dict:
        m = self.fresh_machine()
        lock = None
        if check_lock:
            lock = BR.Session(self.T, self.fresh_machine())
            lock.history(ids)

        x = torch.tensor([ids], device=self.dev, dtype=torch.long)
        core.snapshot(self.c, m)
        n_code = len(self.code_dims)

        if self.executor and not self.verify:
            keeps = self._exec_start_collect(ids)
            cache, logits = self._qwen_prefill(x, keeps, full_head=check_head)
        else:
            lm_out = self.lm(input_ids=x, use_cache=True)
            cache = lm_out.past_key_values
            if self.executor:
                self.erow.start(ids)
            if self.fast:
                ccache = self._chip_cache()
                h = self.cm.embed_tokens(x)
                for L in range(self.n_chip):
                    h = self.cm.layers[L](
                        h,
                        position_embeddings=None,
                        attention_mask=None,
                        position_ids=None,
                        past_key_values=ccache,
                    )
                hf.inject(self.c, ccache)
                if self.graph is None:
                    self._capture()
                g = len(ids)
            elif not self.executor:
                hf.inject(self.c, cache)
            logits = self.model.lm_head(lm_out.last_hidden_state[0, -1])
        pos = len(ids)
        gen: list[int] = []
        n_qwen_text = n_echo_text = n_chip_tok = n_calls = head_mismatch = n_verify_bad = 0
        nxt = int(torch.argmax(logits[: self.v0]))
        prev_owned = False
        in_tool_call = False
        pending_hcs: list[torch.Tensor] = []
        pending_pos0: int | None = None
        t0 = time.time()
        expect: list[int] = []
        lock_error = None

        while n_chip_tok < max_chip_tokens:
            if n_qwen_text >= max_new_text and (not prev_owned) and (not in_tool_call):
                break

            gen.append(nxt)
            if nxt == self.T.TCO and not prev_owned:
                in_tool_call = True
            elif nxt == self.T.TCC:
                in_tool_call = False
            if nxt == self.go_model_id:
                n_calls += 1

            if lock is not None:
                st = nxt - self.v0 if nxt >= self.v0 else BR.text_token(nxt)
                if expect:
                    e = expect.pop(0)
                    if st != e:
                        lock_error = {"at": len(gen) - 1, "got": st, "want": e}
                        break
                elif st < S.VOCAB:
                    lock_error = {"at": len(gen) - 1, "got": S.name(st), "want": "text"}
                    break
                else:
                    k = len(lock.toks)
                    try:
                        lock.qwen([nxt])
                    except RuntimeError as exc:
                        lock_error = {"at": len(gen) - 1, "got": str(exc), "want": "halt"}
                        break
                    expect = list(lock.toks[k + 1 :])

            if nxt in self.eos and not prev_owned:
                break

            chip_in = nxt >= self.v0
            hc = keep = None
            if self.executor:
                keep = self._exec_step_one(nxt)
                if not chip_in:
                    hc = self._exec_hc(nxt, keep)
            if self.fast and (self.verify or not self.executor):
                hf_hc = self._chip_step(nxt, g, ccache)
                g += 1
                if self.verify:
                    same = (
                        bool((hf_hc[self.keep_cols] == keep.to(torch.bfloat16)).all())
                        if chip_in
                        else bool((hf_hc == hc).all())
                    )
                    if not same:
                        n_verify_bad += 1
                hc = hf_hc
            elif not self.executor:
                hc, logits = self._layers_slow(nxt, cache, None if chip_in else pos)

            if chip_in:
                n_chip_tok += 1
                nxt = (
                    self._decode_code_list(keep[:n_code].tolist())
                    if hc is None
                    else self.circuit_token(hc)
                )
                prev_owned = True
                if nxt is None:
                    raise RuntimeError("Circuit relinquished ownership mid-chip-token")
            else:
                if prev_owned:
                    n_echo_text += 1
                else:
                    n_qwen_text += 1
                ct = (
                    self._decode_code_list(keep[:n_code].tolist())
                    if (keep is not None and not self.verify)
                    else self.circuit_token(hc)
                )
                if ct is not None and (self.fast or self.executor) and not check_head:
                    if pending_pos0 is None:
                        pending_pos0 = pos
                    pending_hcs.append(hc)
                    pos += 1
                    nxt = ct
                    prev_owned = True
                else:
                    if self.fast or self.executor:
                        if pending_hcs:
                            pending_hcs.append(hc)
                            logits = self._qwen_chunk(
                                pending_hcs, cache, pending_pos0, full_head=check_head
                            )
                            pending_hcs.clear()
                            pending_pos0 = None
                        else:
                            logits = self._qwen_step(hc, cache, pos, full_head=check_head)
                    pos += 1
                    if ct is not None:
                        if check_head and int(torch.argmax(logits)) != ct:
                            head_mismatch += 1
                        nxt = ct
                        prev_owned = True
                    else:
                        q_logits = logits[: self.v0]
                        if n_calls >= max_calls or n_qwen_text >= max(16, max_new_text - 48):
                            q_logits = q_logits.clone()
                            q_logits[self.T.TCO] = float("-inf")
                        nxt = int(torch.argmax(q_logits))
                        prev_owned = False

        dt = time.time() - t0
        parsed = self.parse_turn(gen)
        return {
            "tokens": gen,
            "text_tokens": n_qwen_text,
            "echo_text_tokens": n_echo_text,
            "chip_tokens": n_chip_tok,
            "tool_calls_executed": n_calls,
            "seconds": round(dt, 3),
            "head_mismatch": head_mismatch,
            "lock_error": lock_error,
            "verify_mismatch": n_verify_bad if self.verify else None,
            **parsed,
        }

    def text_of(self, ids: list[int]) -> str:
        return self.tok.decode([i for i in ids if i < self.v0], skip_special_tokens=False)

    def parse_turn(self, gen_ids: list[int]) -> dict:
        raw = self.text_of(gen_ids)
        calls = re.findall(r"<tool_call>\s*(.*?)\s*</tool_call>", raw, flags=re.DOTALL)
        responses = re.findall(r"<tool_response>\n?(.*?)\n?</tool_response>", raw, flags=re.DOTALL)
        cleaned = re.sub(
            r"<tool_call>.*?</tool_response>\s*(?:<\|im_end\|>\s*<\|im_start\|>assistant\s*<think>\s*</think>\s*)?",
            "\n\n",
            raw,
            flags=re.DOTALL,
        )
        cleaned = re.sub(r"<tool_call>.*$", "", cleaned, flags=re.DOTALL)
        cleaned = re.sub(r"</?think>", "", cleaned)
        for stop_s in ("<|im_end|>", "<|endoftext|>", "<|im_start|>"):
            if stop_s in cleaned:
                cleaned = cleaned.split(stop_s, 1)[0]
        cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip()
        return {
            "raw_text": raw,
            "tool_calls": calls,
            "tool_responses": responses,
            "final_explanation": cleaned,
        }


def main() -> None:
    ap = argparse.ArgumentParser(description="Run Qwen3.5-Math-GraphCalc in-weights")
    ap.add_argument(
        "--ask",
        default="Graph y = x^3 - 3*x from x = -2 to 2, identify its roots and turning points, and state the local maximum y-value.",
        help="Math or graphing question to solve",
    )
    ap.add_argument("--pure-hf", action="store_true", help="Run un-fused HF Qwen3_5DecoderLayer weights (executor=False)")
    ap.add_argument("--max-new-text", type=int, default=512)
    args = ap.parse_args()

    runner = CarbonGraphCalcRunner(MODEL_DIR, fast=True, executor=not args.pure_hf)
    res = runner.generate(runner.prompt_ids(args.ask), max_new_text=args.max_new_text, check_lock=True)
    if res["tool_responses"]:
        print("=== IN-WEIGHTS RV32 <tool_response> ===")
        print(res["tool_responses"][0])
        print()
    print("=== STEP-BY-STEP TUTOR EXPLANATION ===")
    print(res["final_explanation"])
    print(f"\n[chip_tokens={res['chip_tokens']} | lock_error={res['lock_error']} | {res['seconds']}s]")


if __name__ == "__main__":
    main()