#!/usr/bin/env python
"""K5: inference-time levers at MATCHED total generated tokens per problem (SOTA plan v3.1 section 3; docs/K5_LEVERS.md).

Three arms on the same base model (Goedel-Prover-V2-8B) and the same dev problems:

  restart   independent whole-proof attempts (class A samples).
  correct   Goedel-format Lean-feedback correction (class B): attempt + up to 2 repairs per trajectory, cumulative
            token ledger, repairs of EVERY failed root (so each trajectory has a complete cost record).
  agentic   (class C, exploratory) proof-bank-conditioned attempt, then R = 4 rounds in total of: APOLLO-style subgoal
            repair (split the failing proof at its failing `have`s, prove each subgoal as its own problem, reassemble,
            re-verify) alternating with compiler-feedback correction. Retrieval is from a decontaminated bank of
            verified proofs of OTHER training problems.

Matching. The unit of every arm is a TRAJECTORY (a restart sample; a root with its repair chain; an agentic chain with
its subgoal sub-generations) with a token cost: the completion tokens spent up to the end of the round in which it first
succeeded, or all of its tokens if it never did. Arms are generated generously (a pool of units per problem) and then
REPLAYED at a per-problem budget B_p: units are drawn in a random order (permutation-averaged, deterministic seed) and
run until the budget is spent or one succeeds; a unit that does not fit in the remaining budget ends the run (a real
run cannot know a unit's cost in advance, so it is clipped and fails). The estimate per problem is the fraction of
orderings that solve it, directly comparable across arms and testable with the registered paired sign-flip test.
`--budget restart_k:8` sets B_p to 8 mean restart attempts of that problem (the pass@8-equivalent spend);
`--budget fixed:N` is a scalar. A problem whose pool cannot fill the budget is reported `underfilled`, and
`status` says how many more units to generate: the budget equaliser never silently lets one arm spend less.

Subcommands (`k5_levers.py <cmd> -h`):
  plan      print every stage (the exact command lines) for a configuration; starts nothing
  run       execute the stages with a resumable done-marker per stage (each inside `hostguard.py run`)
  prep      merge dev problems, build + decontaminate the proof bank, attach retrieved examples (CPU)
  route     agentic round k: choose subgoal vs feedback per failed parent, extract subgoals through Lean (CPU)
  assemble  agentic round k: splice proven subgoals back, write assembled candidates (CPU)
  ledger    per-arm unit tables from generation + verification files
  report    matched-budget comparison, paired tests, Holm over the K5 family
  status    under-filled problems per arm and how many units each needs
  estimate  box-hour estimate from measured rates
"""

from __future__ import annotations

import argparse
import dataclasses
import hashlib
import importlib.util
import json
import math
import os
import re
import shlex
import subprocess
import sys
import time
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Optional

import numpy as np

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))


def _load(name: str, rel: str):
    if name in sys.modules:
        return sys.modules[name]
    spec = importlib.util.spec_from_file_location(name, REPO / rel)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


C = _load("p2_common", "scripts/p2/p2_common.py")
LP = _load("k5_lean_prompts", "extras/vllm_rollout/lean_prompts.py")

# a git worktree has no venvs of its own: ALIGN_VENV_ROOT (default: the main checkout) holds .venv and .venv-rollout
_VR = Path(os.environ.get("ALIGN_VENV_ROOT", REPO if (REPO / ".venv").exists() else "/apps/development/rlhf/align"))
PYR = str(_VR / ".venv-rollout" / "bin" / "python")
PY = str(_VR / ".venv" / "bin" / "python")
DEFAULT_MODEL = (
    "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
    "dfd02e6271a58375dfbf3ece0175277cf6b6a89a"
)
ARMS = ("restart", "correct", "agentic")
ASM_BASE = 1_000_000  # sample_idx offset of assembled (subgoal-repaired) nodes: never collides with lean_gen's parent*n+i


# =========================================================================== 1. ledger and the budget equaliser


@dataclass(frozen=True)
class Unit:
    """One trajectory of an arm. `tokens` is the cost the unit charges when it is run: the completion tokens of every
    node up to the end of the round of its first success, or `tokens_total` when it never succeeded."""

    problem_id: str
    arm: str
    unit_id: str
    tokens: int
    solved: bool
    tokens_total: int
    by_kind: tuple = ()  # ((kind, tokens), ...) over the charged span


def node_tokens(g: dict) -> int:
    t = g.get("trajectory")
    if t and "completion_tokens" in t:
        return int(t["completion_tokens"])
    ids = g.get("completion_ids")
    if ids is None:
        raise ValueError(f"{g.get('problem_id')}: generation record without completion_ids or a trajectory ledger")
    return len(ids)


def node_kind(g: dict) -> str:
    t = g.get("trajectory") or {}
    return t.get("kind") or ("attempt" if int(g["round"]) == 0 else "repair")


def root_id(g: dict) -> str:
    t = g.get("trajectory")
    if t and "root_trajectory_id" in t:
        return t["root_trajectory_id"]
    if int(g["round"]) != 0:
        raise ValueError(f"{g['problem_id']}: a round-{g['round']} record without a trajectory ledger cannot be assigned to a root")
    return f"{g['problem_id']}|{int(g['sample_idx'])}"


def _join(gens: Iterable[dict], vers: Iterable[dict]) -> list[tuple[dict, dict]]:
    vmap = {}
    for v in vers:
        vmap[C.gen_key(v)] = v
    out, seen = [], set()
    for g in gens:
        k = C.gen_key(g)
        if k in seen:
            raise ValueError(f"duplicate generation key {k} across the input files")
        seen.add(k)
        v = vmap.get(k)
        if v is None:
            raise ValueError(f"generation {k} has no verification record; verify every node before building the ledger")
        out.append((g, v))
    return out


def units_restart(gens: Iterable[dict], vers: Iterable[dict], arm: str = "restart") -> dict[str, list[Unit]]:
    """Each round-0 sample is its own unit."""
    out: dict[str, list[Unit]] = defaultdict(list)
    for g, v in _join(gens, vers):
        if int(g["round"]) != 0:
            continue
        n = node_tokens(g)
        out[g["problem_id"]].append(Unit(g["problem_id"], arm, root_id(g), n, bool(v["accepted"]), n, ((node_kind(g), n),)))
    return dict(out)


def units_chains(gens: Iterable[dict], vers: Iterable[dict], arm: str) -> dict[str, list[Unit]]:
    """One unit per root trajectory (class B and agentic). Nodes are grouped by `trajectory.root_trajectory_id`; the
    cost through round r is the sum of ALL nodes of rounds <= r (branches included), so nothing generated is free."""
    groups: dict[str, list[tuple[dict, dict]]] = defaultdict(list)
    for g, v in _join(gens, vers):
        groups[root_id(g)].append((g, v))
    out: dict[str, list[Unit]] = defaultdict(list)
    for rid, nodes in groups.items():
        if not any(int(g["round"]) == 0 for g, _ in nodes):
            raise ValueError(f"{rid}: no round-0 node among the input files (the lineage is partial); pass every round's files")
        pid = nodes[0][0]["problem_id"]
        solved_round = min((int(g["round"]) for g, v in nodes if v["accepted"]), default=None)
        charged = [(g, v) for g, v in nodes if solved_round is None or int(g["round"]) <= solved_round]
        by_kind: Counter = Counter()
        for g, _ in charged:
            by_kind[node_kind(g)] += node_tokens(g)
        total = sum(node_tokens(g) for g, _ in nodes)
        out[pid].append(Unit(pid, arm, rid, sum(by_kind.values()), solved_round is not None, total, tuple(sorted(by_kind.items()))))
    return dict(out)


def _seed_for(seed: int, pid: str) -> int:
    return int.from_bytes(hashlib.sha256(f"{seed}|{pid}".encode()).digest()[:8], "big") % (2**32)


@dataclass
class Replay:
    p_solve: float  # fraction of orderings in which the problem is solved within the budget
    spent: float  # mean tokens spent over the orderings
    underfill: float  # fraction of orderings that ran out of units with a mean unit of budget still unspent
    n_units: int


def replay_budget(units: list[Unit], budget: float, n_perm: int = 256, seed: int = 0, pid: str = "") -> Replay:
    """Permutation-averaged solve probability of `units` at `budget` generated tokens (the equaliser's arithmetic).

    One ordering: units are run in order; a unit whose cost exceeds the remaining budget ends the ordering unsolved
    (it would be clipped); a solved unit within budget ends it solved; an unsolved one that fits spends its tokens
    and the next is tried. Spent tokens are what the ordering actually used."""
    if not units:
        return Replay(0.0, 0.0, 1.0, 0)
    rng = np.random.default_rng(_seed_for(seed, pid))
    cost = np.array([u.tokens for u in units], dtype=np.float64)
    solved = np.array([u.solved for u in units], dtype=bool)
    mean_cost = float(cost.mean())
    wins = under = 0
    spent_sum = 0.0
    n = len(units)
    for _ in range(n_perm):
        order = rng.permutation(n)
        spent, won, exhausted = 0.0, False, True
        for i in order:
            if spent + cost[i] > budget:
                exhausted = False
                break
            spent += cost[i]
            if solved[i]:
                won, exhausted = True, False
                break
        wins += won
        spent_sum += spent
        if exhausted and not won and budget - spent >= mean_cost:
            under += 1
    return Replay(wins / n_perm, spent_sum / n_perm, under / n_perm, n)


def budget_table(spec: str, restart: dict[str, list[Unit]], problems: Iterable[str]) -> dict[str, float]:
    """Per-problem budget B_p. `restart_k:K` = K x the mean tokens of that problem's restart attempts (the spend of a
    pass@K run); `fixed:N` = N tokens everywhere."""
    kind, _, arg = spec.partition(":")
    out = {}
    if kind == "fixed":
        return {p: float(arg) for p in problems}
    if kind != "restart_k":
        raise ValueError(f"unknown budget spec {spec!r} (restart_k:K or fixed:N)")
    k = float(arg)
    for p in problems:
        us = restart.get(p)
        if not us:
            raise ValueError(f"problem {p}: no restart units to derive its budget from")
        out[p] = k * sum(u.tokens_total for u in us) / len(us)
    return out


def equalise(arms: dict[str, dict[str, list[Unit]]], budgets: dict[str, float], n_perm: int = 256, seed: int = 0,
             fractions: tuple = (0.25, 0.5, 1.0)) -> dict:
    """Replay every arm at every budget fraction. -> {arm: {fraction: {problem: Replay}}} plus summaries."""
    res: dict = {}
    for arm, by_p in arms.items():
        res[arm] = {}
        for f in fractions:
            res[arm][f] = {p: replay_budget(by_p.get(p, []), b * f, n_perm, seed, p) for p, b in budgets.items()}
    return res


def arm_summary(per_problem: dict[str, Replay], budgets: dict[str, float], fraction: float = 1.0) -> dict:
    n = len(per_problem)
    tot_b = sum(budgets[p] * fraction for p in per_problem)
    return {
        "n_problems": n,
        "mean_p_solve": sum(r.p_solve for r in per_problem.values()) / max(n, 1),
        "tokens_spent_per_problem": sum(r.spent for r in per_problem.values()) / max(n, 1),
        "budget_per_problem": tot_b / max(n, 1),
        "spend_ratio": sum(r.spent for r in per_problem.values()) / tot_b if tot_b else 0.0,
        "underfilled_problems": sum(1 for r in per_problem.values() if r.underfill > 0.05),
        "no_units": sum(1 for r in per_problem.values() if r.n_units == 0),
    }


def units_needed(units: dict[str, list[Unit]], budgets: dict[str, float], per_problem: dict[str, Replay]) -> dict[str, int]:
    """Additional units per under-filled problem: enough mean-cost units to cover the unspent budget, +1 slack."""
    need = {}
    for p, r in per_problem.items():
        if r.underfill <= 0.05:
            continue
        us = units.get(p) or []
        mean = (sum(u.tokens for u in us) / len(us)) if us else max(budgets[p] / 4.0, 1.0)
        need[p] = max(1, math.ceil(max(budgets[p] - r.spent, 0.0) / max(mean, 1.0))) + 1
    return need


# =========================================================================== 2. proof bank + decontamination


@dataclass
class BankEntry:
    problem_id: str
    statement: str  # `theorem ... := by` (no header, no `sorry`)
    proof: str  # text after that `:= by`, verbatim
    informal: Optional[str]
    tokens: int
    code_sha: str
    source_statement: str = ""  # the full formal statement (for decontamination)


def split_example(lean_code: str) -> Optional[tuple[str, str]]:
    """(`theorem ... := by`, body) of a verified `lean_code`, or None. Comments are masked for the search only."""
    from align.rl.envs.subgoal_repair import mask

    mk = mask(lean_code)
    m = re.search(r"\btheorem\b[\s\S]*?:=\s*by", mk)
    if m is None:
        return None
    return lean_code[m.start():m.end()], lean_code[m.end():]


def eval_items_from_problems(rows: Iterable[dict]):
    from align.data.decontam import EvalItem

    return [EvalItem(f"{r.get('split', '')}:{r['problem_id']}", r.get("source", "eval"), r["formal_statement"], r.get("informal_prefix"))
            for r in rows]


def decontaminate_bank(candidates: list[BankEntry], eval_rows: list[dict], ancestry: Optional[dict] = None) -> tuple[list[BankEntry], list[dict]]:
    """Drop every candidate that hits any evaluation / query problem under the repo's layered detector
    (`align.data.decontam`: normalised-text, signature and alpha hashes, MinHash-LSH, token n-gram containment, informal
    text n-grams), or whose problem id IS an evaluation id, or whose ancestry (mutation parent) does. Returns
    (kept, dropped-with-reason). Over-flagging is the policy: a false positive costs one bank entry."""
    from align.data.decontam import ContaminationLedger, Decontaminator

    dec = Decontaminator(eval_items_from_problems(eval_rows))
    eval_ids = {r["problem_id"] for r in eval_rows}
    led = ContaminationLedger()
    for c in candidates:
        led.add(c.problem_id, dec.check(c.source_statement or c.statement, c.informal), (ancestry or {}).get(c.problem_id))
    kept, dropped = [], []
    for c in candidates:
        if c.problem_id in eval_ids:
            dropped.append({"problem_id": c.problem_id, "why": "is_an_eval_problem"})
            continue
        bad, why = led.resolve(c.problem_id)
        if bad:
            dropped.append({"problem_id": c.problem_id, "why": why, "hits": [dataclasses.asdict(h) for h in led.direct.get(c.problem_id, [])[:3]]})
            continue
        kept.append(c)
    return kept, dropped


def bank_candidates(problems: dict[str, dict], gens: Iterable[dict], vers: Iterable[dict], max_chars: int = 2500) -> list[BankEntry]:
    """The shortest accepted proof per problem, as an example (statement + body)."""
    best: dict[str, BankEntry] = {}
    for g, v in _join(gens, vers):
        if not v["accepted"] or not g.get("lean_code") or g["problem_id"] not in problems:
            continue
        ex = split_example(g["lean_code"])
        if ex is None or len(ex[0]) + len(ex[1]) > max_chars:
            continue
        pr = problems[g["problem_id"]]
        e = BankEntry(g["problem_id"], ex[0], ex[1], pr.get("informal_prefix"), node_tokens(g), C.code_sha(g["lean_code"]), pr["formal_statement"])
        cur = best.get(g["problem_id"])
        if cur is None or len(e.proof) < len(cur.proof):
            best[g["problem_id"]] = e
    return [best[k] for k in sorted(best)]


class BankIndex:
    """TF-IDF over statement tokens (identifiers kept: lemma names are the signal). Cosine retrieval, shorter proof breaks
    ties. Pure numpy; the bank is thousands of entries."""

    def __init__(self, entries: list[BankEntry]):
        from align.data.decontam import lean_tokens

        self._tok = lambda s: lean_tokens(s, alpha=False)
        self.entries = entries
        docs = [Counter(self._tok(e.source_statement or e.statement)) for e in entries]
        df = Counter(t for d in docs for t in d)
        self.vocab = {t: i for i, t in enumerate(sorted(df))}
        n = max(len(entries), 1)
        self.idf = np.array([math.log((1 + n) / (1 + df[t])) + 1.0 for t in sorted(df)]) if df else np.zeros(0)
        self.mat = np.zeros((len(entries), len(self.vocab)), dtype=np.float32)
        for i, d in enumerate(docs):
            self.mat[i] = self._vec(d)

    def _vec(self, counts: Counter) -> np.ndarray:
        v = np.zeros(len(self.vocab), dtype=np.float32)
        for t, c in counts.items():
            j = self.vocab.get(t)
            if j is not None:
                v[j] = (1 + math.log(c)) * self.idf[j]
        nrm = np.linalg.norm(v)
        return v / nrm if nrm else v

    def retrieve(self, statement: str, k: int = 3, exclude_ids: Iterable[str] = (), min_sim: float = 0.05) -> list[BankEntry]:
        if not self.entries:
            return []
        q = self._vec(Counter(self._tok(statement)))
        sims = self.mat @ q
        ex = set(exclude_ids)
        order = sorted(range(len(self.entries)), key=lambda i: (-float(sims[i]), len(self.entries[i].proof)))
        out = []
        for i in order:
            if sims[i] < min_sim or self.entries[i].problem_id in ex:
                continue
            out.append(self.entries[i])
            if len(out) == k:
                break
        return out


def bank_examples(entries: list[BankEntry], max_chars: int = 6000) -> list[dict]:
    out, used = [], 0
    for e in entries:
        used += len(e.statement) + len(e.proof)
        if used > max_chars:
            break
        out.append({"statement": e.statement, "proof": e.proof, "source": e.problem_id})
    return out


def attach_bank(problems: list[dict], index: BankIndex, k: int = 3) -> list[dict]:
    """Problem rows with `bank_examples` (the `goedel_bank` prompt format reads them). A problem is never given its own
    bank entry (`exclude_ids`)."""
    out = []
    for p in problems:
        ex = bank_examples(index.retrieve(p["formal_statement"], k, exclude_ids=[p["problem_id"]]))
        out.append(dict(p, bank_examples=ex))
    return out


# =========================================================================== 3. agentic rounds (route / assemble)


def parent_info(g: dict, v: dict) -> Optional[dict]:
    """Failed node with Lean feedback, or None (accepted, truncated, no code, no compile errors)."""
    if v["accepted"] or not g.get("lean_code"):
        return None
    if not isinstance(v.get("repl_errors"), list) or not v["repl_errors"] or not isinstance(v.get("repl_code"), str):
        return None
    return {"gen": g, "ver": v}


def route_for_round(routes: str, round_k: int) -> str:
    """`routes` = comma list for rounds 1..R-1, each `subgoal` (subgoal repair, falling back to feedback when nothing can
    be extracted) or `feedback`. A design choice, not tuned: alternating heterogeneous moves."""
    rs = [x.strip() for x in routes.split(",")]
    if not 1 <= round_k <= len(rs) or any(x not in ("subgoal", "feedback") for x in rs):
        raise ValueError(f"routes {routes!r} has no valid entry for round {round_k}")
    return rs[round_k - 1]


def sg_problem_id(parent_g: dict, j: int) -> str:
    return f"{parent_g['problem_id']}::sg{int(parent_g['sample_idx'])}r{int(parent_g['round'])}g{j}"


def make_asm_nodes(plan, parent_g: dict, parent_v: dict, bodies: dict[int, str], sub_tokens: int, n_proven: int, round_k: int, cap: int):
    """(gen record, ver record or None). The assembled candidate is a synthetic generation: its text is the code, its
    completion_ids a single placeholder, and the real cost (the subgoal samples) sits in `trajectory.completion_tokens`.
    The history `messages` are the parent's, so a following Goedel correction prompt reads the whole story. When the
    candidate equals the parent's code (no subgoal proven) its verdict is copied, not re-verified."""
    from align.rl.envs.subgoal_repair import reassemble

    code = reassemble(plan, bodies, restore_unproven=True)
    pt = parent_g.get("trajectory") or {}
    prior = int(pt.get("cum_completion_tokens", node_tokens(parent_g)))
    sidx = ASM_BASE + int(parent_g["sample_idx"])
    rid = root_id(parent_g)
    gen = {
        "problem_id": parent_g["problem_id"], "sample_idx": sidx, "round": round_k, "cap_tokens": cap,
        "prompt_ids": [0], "completion_ids": [0], "prompt_text": "", "completion_text": f"```lean4\n{code}\n```",
        "finish_reason": "stop", "lean_code": code, "lean_code_error": None, "messages": parent_g["messages"],
        "model_path": parent_g["model_path"], "model_rev": parent_g["model_rev"],
        "sampling": dict(parent_g["sampling"], max_tokens=1),
        "context_clipped": False,
        "parent": {"problem_id": parent_g["problem_id"], "sample_idx": int(parent_g["sample_idx"]), "round": int(parent_g["round"])},
        "trajectory": {
            "root_trajectory_id": rid, "problem_id": parent_g["problem_id"],
            "root_sample_idx": int(pt.get("root_sample_idx", parent_g["sample_idx"])), "round": round_k, "kind": "assembled",
            "parent_key": [parent_g["problem_id"], int(parent_g["sample_idx"]), int(parent_g["round"])],
            "completion_tokens": int(sub_tokens), "prior_completion_tokens": prior, "cum_completion_tokens": prior + int(sub_tokens),
            "truncated": False, "budget_clipped": False, "partial": False,
        },
        "subgoal": {"gaps": len(plan.gaps), "proven": n_proven},
    }
    ver = None
    if code == parent_g["lean_code"]:
        ver = dict(parent_v, sample_idx=sidx, round=round_k, cap_tokens=cap, code_sha256=C.code_sha(code))
    return gen, ver


# =========================================================================== 4. stages (the exact command lines)


@dataclass
class Cfg:
    run_dir: str
    problems: list[str]  # dev problem files (miniF2F-valid, ProofNet-valid, ...)
    model: str = DEFAULT_MODEL
    seed: int = 0
    cap: int = 40960
    traj_cap: int = 40000  # class B cumulative cap
    n_restart: int = 32
    n_agentic: int = 8
    rounds: int = 4  # agentic rounds in total (1 attempt + 3 repairs)
    routes: str = "subgoal,feedback,subgoal"
    sub_n: int = 4  # samples per subgoal
    sub_cap: int = 16384
    bank_k: int = 3
    budget: str = "restart_k:8"
    workers: int = 8
    gpu_mem_util: float = 0.82
    bank_problems: Optional[str] = None
    bank_gens: Optional[str] = None
    bank_vers: Optional[str] = None
    eval_problems: list = field(default_factory=list)  # every evaluation set the bank must be clean against
    hostguard: bool = True
    gpu_caps: str = "--cores 24 --mem-gb 100"
    cpu_caps: str = "--cores 16 --mem-gb 64"

    def p(self, *parts) -> str:
        return str(Path(self.run_dir, *parts))


@dataclass
class Stage:
    name: str
    argv: list
    gpu: bool
    outputs: list  # a stage is done iff its marker exists; outputs are listed for `plan` and cleanup
    skip_if_empty: Optional[str] = None  # an input file that, when it exists with no rows, makes the stage a no-op

    def marker(self, cfg: Cfg) -> Path:
        return Path(cfg.p("done", self.name))


def _hg(cfg: Cfg, st: Stage) -> list:
    if not cfg.hostguard:
        return st.argv
    caps = (cfg.gpu_caps if st.gpu else cfg.cpu_caps).split()
    return [PY, str(REPO / "scripts/hostguard.py"), "run", "--wait", *caps, "--label", f"k5-{st.name}", "--", *st.argv]


def _gen_args(cfg: Cfg, problems: str, out: str, n: int, rnd: int, cap: int, label: str, extra: list = ()) -> str:
    a = ["--model", cfg.model, "--problems", problems, "--out", out, "--n", str(n), "--cap-tokens", str(cap), "--round", str(rnd),
         "--tp", "4", "--seed", str(cfg.seed), "--max-num-seqs", "128", "--gpu-mem-util", str(cfg.gpu_mem_util), "--label", label, *extra]
    return " ".join(shlex.quote(x) for x in a)


def _pgv(cfg: Cfg, name: str, problems: str, gen_out: str, ver_out: str, n: int, rnd: int, cap: int, extra: list = (), feedback: bool = True,
         wait_marker: Optional[str] = None) -> Stage:
    run = cfg.p("pipe", name)
    argv = [PY, str(REPO / "scripts/p2/pipeline_gen_verify.py"), "--run-dir", run, "--problems", problems, "--verify-out", ver_out,
            "--workers", str(cfg.workers), "--label", f"k5-{name}", "--gen-args", _gen_args(cfg, problems, gen_out, n, rnd, cap, f"k5-{name}", list(extra))]
    if feedback:
        argv.append("--feedback")
    return Stage(name, argv, True, [gen_out, ver_out])


def _self(cfg: Cfg, name: str, sub: str, *args: str, gpu: bool = False) -> Stage:
    return Stage(name, [PY, str(REPO / "scripts/p2/k5_levers.py"), sub, *args], gpu, [])


def build_stages(cfg: Cfg) -> list[Stage]:
    """Every stage of K5, in order. Pure: nothing is executed or created."""
    dev, dev_bank = cfg.p("dev.jsonl"), cfg.p("dev_bank.jsonl")
    ss: list[Stage] = []
    prep = ["--run-dir", cfg.run_dir, "--problems", *cfg.problems, "--bank-k", str(cfg.bank_k)]
    if cfg.bank_problems:
        prep += ["--bank-problems", cfg.bank_problems, "--bank-gens", cfg.bank_gens, "--bank-vers", cfg.bank_vers]
    if cfg.eval_problems:
        prep += ["--eval-problems", *cfg.eval_problems]
    ss.append(_self(cfg, "prep", "prep", *prep))
    # ---- restart (class A samples; also the roots of the correction arm)
    R0g, R0v = cfg.p("restart/gen_r0.jsonl"), cfg.p("restart/ver_r0.jsonl")
    ss.append(_pgv(cfg, "restart_r0", dev, R0g, R0v, cfg.n_restart, 0, cfg.cap))
    # ---- correct (class B): rounds 1, 2; every failed root is repaired (--repair-solved) so each trajectory is fully costed
    prev = [(R0v, R0g)]
    for r in (1, 2):
        g, v = cfg.p(f"correct/gen_r{r}.jsonl"), cfg.p(f"correct/ver_r{r}.jsonl")
        extra = ["--trajectory-cap-tokens", str(cfg.traj_cap), "--repair-solved"]
        for pv, pg in prev:
            extra += ["--corrections-from", pv, pg]
        ss.append(_pgv(cfg, f"correct_r{r}", dev, g, v, 1, r, cfg.cap, extra))
        prev = [(v, g)] + prev
    # ---- agentic (class C)
    A0g, A0v = cfg.p("agentic/gen_r0.jsonl"), cfg.p("agentic/ver_r0.jsonl")
    ss.append(_pgv(cfg, "agentic_r0", dev_bank, A0g, A0v, cfg.n_agentic, 0, cfg.cap, ["--prompt-format", "goedel_bank"]))
    last = [(A0v, A0g)]  # the (ver, gen) pairs of the previous round
    for r in range(1, cfg.rounds):
        mode = route_for_round(cfg.routes, r)
        rd = cfg.p(f"agentic/round{r}")
        route_args = ["--run-dir", cfg.run_dir, "--round", str(r), "--mode", mode, "--workers", "2", "--bank-k", str(cfg.bank_k), "--sub-n", str(cfg.sub_n)]
        for v, g in last:
            route_args += ["--pairs", v, g]
        ss.append(_self(cfg, f"agentic_route_r{r}", "route", *route_args))
        fg, fv = cfg.p(f"agentic/gen_r{r}.jsonl"), cfg.p(f"agentic/ver_r{r}.jsonl")
        fbg, fbv = f"{rd}/fb_parents.gen.jsonl", f"{rd}/fb_parents.ver.jsonl"
        nxt = [(fv, fg)]
        if mode == "subgoal":
            sgp, sgg, sgv = f"{rd}/sg_problems.jsonl", f"{rd}/sg_gen.jsonl", f"{rd}/sg_ver.jsonl"
            st = _pgv(cfg, f"agentic_sg_r{r}", sgp, sgg, sgv, cfg.sub_n, 0, cfg.sub_cap, ["--prompt-format", "goedel_bank"], feedback=False)
            st.skip_if_empty = sgp
            ss.append(st)
        extra = ["--trajectory-cap-tokens", str(cfg.traj_cap), "--repair-solved", "--corrections-from", fbv, fbg]
        st = _pgv(cfg, f"agentic_fb_r{r}", dev_bank, fg, fv, 1, r, cfg.cap, extra)
        st.skip_if_empty = fbg
        ss.append(st)
        if mode == "subgoal":
            ag, av = f"{rd}/asm.gen.jsonl", f"{rd}/asm.ver.jsonl"
            ss.append(_self(cfg, f"agentic_asm_r{r}", "assemble", "--run-dir", cfg.run_dir, "--round", str(r), "--cap", str(cfg.cap)))
            st = Stage(f"agentic_asmverify_r{r}", [PY, str(REPO / "scripts/p2/verify_generations.py"), "--generations", ag, "--problems", dev,
                                                      "--out", av, "--workers", str(cfg.workers), "--budget-s", "300", "--feedback"], False, [av])
            st.skip_if_empty = ag
            ss.append(st)
            nxt.append((av, ag))
        last = nxt
    ss.append(_self(cfg, "report", "report", "--run-dir", cfg.run_dir, "--budget", cfg.budget, "--rounds", str(cfg.rounds)))
    return ss


class Runner:
    def run(self, argv: list, log: Path) -> int:
        log.parent.mkdir(parents=True, exist_ok=True)
        with open(log, "ab") as lf:
            return subprocess.call(argv, stdout=lf, stderr=subprocess.STDOUT, cwd=str(REPO))


def run_stages(cfg: Cfg, stages: list[Stage], runner, say=print, only: Optional[set] = None) -> int:
    for st in stages:
        if only and st.name not in only:
            continue
        mk = st.marker(cfg)
        if mk.exists():
            say(f"skip {st.name} (done)")
            continue
        if st.skip_if_empty and Path(st.skip_if_empty).exists() and not Path(st.skip_if_empty).read_text().strip():
            say(f"skip {st.name} (its input {st.skip_if_empty} is empty)")
            mk.parent.mkdir(parents=True, exist_ok=True)
            mk.write_text("empty")
            continue
        say(f"START {st.name}")
        rc = runner.run(_hg(cfg, st), Path(cfg.p("logs", st.name + ".log")))
        say(f"END {st.name} rc={rc}")
        if rc != 0:
            say(f"STOP: {st.name} failed; re-running this command resumes (every stage is keyed and resumable)")
            return rc
        mk.parent.mkdir(parents=True, exist_ok=True)
        mk.write_text(time.strftime("%F %T"))
    return 0


# =========================================================================== 5. estimate


MEASURED = {  # docs/K5_LEVERS.md section 7: from the P2 baseline manifest (Goedel-8B, miniF2F-test, cap 40960)
    "box_tok_per_h": 522.5 * 3600,  # end-to-end tokens/s of the whole 4xA10G box, generation overlapped with Lean
    "mean_attempt_tokens": {"minif2f": 5685.0, "proofnet": 9000.0},  # proofnet: ASSUMED 1.6x (no measurement yet)
    "fail_rate": {"minif2f": 0.36, "proofnet": 0.75},  # 1 - single-sample pass rate; proofnet ASSUMED
}


def estimate(sets: dict[str, int], cfg: Cfg, m: dict = MEASURED, repair_cost: float = 1.15, repair_success: float = 0.12,
             sub_share: float = 0.35) -> dict:
    """Box-hour estimate per arm. `sets` = {"minif2f": N, "proofnet": N}. The repair multiplier covers the longer
    (history-carrying) prompts; sub_share is the tokens a subgoal round spends relative to one attempt per failed parent."""
    out: dict = {"assumptions": {"repair_cost": repair_cost, "repair_success": repair_success, "sub_share": sub_share, **{"box_tok_per_h": m["box_tok_per_h"]}}}
    tot: Counter = Counter()
    for s, n in sets.items():
        a, f = m["mean_attempt_tokens"][s], m["fail_rate"][s]
        restart = n * cfg.n_restart * a
        f1 = f
        f2 = f * (1 - repair_success)
        correct = n * cfg.n_restart * a * repair_cost * (f1 + f2)
        agent_root = n * cfg.n_agentic * a
        agent = agent_root
        fa = f * 0.85  # bank-conditioned roots assumed slightly better
        for r in range(1, cfg.rounds):
            mode = route_for_round(cfg.routes, r)
            k = n * cfg.n_agentic * fa
            agent += k * a * (sub_share if mode == "subgoal" else repair_cost)
            fa *= (1 - repair_success)
        tot["restart"] += restart
        tot["correct"] += correct
        tot["agentic"] += agent
    out["tokens"] = dict(tot)
    out["box_hours"] = {k: v / m["box_tok_per_h"] for k, v in tot.items()}
    out["box_hours"]["total"] = sum(out["box_hours"].values())
    return out


# =========================================================================== 6. CLI


def _read(path) -> list[dict]:
    return list(C.iter_jsonl(path)) if path and Path(path).exists() else []


def _problems(files: Iterable[str]) -> dict[str, dict]:
    out: dict[str, dict] = {}
    for f in files:
        for pid, rec in C.load_problems(f).items():
            if pid in out and out[pid]["formal_statement"] != rec["formal_statement"]:
                raise C.ContractError(f"problem {pid!r} appears in two files with different statements")
            out[pid] = rec
    return out


def _write(path, rows) -> int:
    return C.write_jsonl_atomic(path, rows)


def cmd_prep(a) -> int:
    rd = Path(a.run_dir)
    probs = _problems(a.problems)
    _write(rd / "dev.jsonl", probs.values())
    bank_entries: list[BankEntry] = []
    report: dict = {"dev_problems": len(probs)}
    if a.bank_problems:
        train = _problems([a.bank_problems])
        cands = bank_candidates(train, _read(a.bank_gens), _read(a.bank_vers))
        eval_rows = list(probs.values()) + [r for f in a.eval_problems for r in _read(f)]
        bank_entries, dropped = decontaminate_bank(cands, eval_rows)
        report.update(bank_candidates=len(cands), bank_kept=len(bank_entries), bank_dropped=len(dropped),
                      eval_items=len(eval_rows), dropped_reasons=dict(Counter(d["why"].split(":")[0] for d in dropped)))
        _write(rd / "bank.jsonl", (dataclasses.asdict(e) for e in bank_entries))
        _write(rd / "bank_dropped.jsonl", dropped)
    rows = attach_bank(list(probs.values()), BankIndex(bank_entries), a.bank_k)
    _write(rd / "dev_bank.jsonl", rows)
    report["with_examples"] = sum(1 for r in rows if r["bank_examples"])
    (rd / "prep.json").write_text(json.dumps(report, indent=1))
    print(json.dumps(report))
    return 0


def _pairs(args) -> tuple[list[dict], list[dict]]:
    gens, vers = [], []
    for v, g in args:
        gens += _read(g)
        vers += _read(v)
    return gens, vers


def cmd_route(a) -> int:
    """Agentic round k. Parents = failed, feedback-bearing nodes of round k-1 (every `--pairs` file). Subgoal rounds
    plan each parent through Lean; parents with nothing extractable, and every parent of a feedback round, go to the
    Goedel correction path via filtered copies of their (ver, gen) records."""
    from align.rl.envs import subgoal_repair as SR

    rd = Path(a.run_dir, "agentic", f"round{a.round}")
    rd.mkdir(parents=True, exist_ok=True)
    probs = {r["problem_id"]: r for r in _read(Path(a.run_dir, "dev_bank.jsonl"))}
    gens, vers = _pairs(a.pairs)
    nodes = [(g, v) for g, v in _join(gens, vers) if int(g["round"]) == a.round - 1]
    st = Counter(parents_total=len(nodes))
    fb, plans, sg_rows = [], [], []
    probe_pool = None
    index = None
    bank = _read(Path(a.run_dir, "bank.jsonl"))
    if bank:
        index = BankIndex([BankEntry(**b) for b in bank])
    for g, v in nodes:
        info = parent_info(g, v)
        if info is None:
            st["no_feedback_or_solved"] += 1
            continue
        route = "feedback"
        if a.mode == "subgoal":
            if probe_pool is None:
                probe_pool = _open_pool(a.workers)
            plan = _plan_with_pool(probe_pool, probs[g["problem_id"]]["formal_statement"], g, v, SR)
            if plan is not None and plan.lemmas:
                route = "subgoal"
                rows = []
                for j, gap in enumerate(plan.gaps):
                    if gap.lemma is None:
                        continue
                    stmt = SR.lemma_problem_statement(plan, gap)
                    row = {"problem_id": sg_problem_id(g, j), "source": "k5_subgoal", "split": "dev", "formal_statement": stmt,
                           "informal_prefix": None, "statement_sha": C.statement_sha(stmt), "parent": [g["problem_id"], int(g["sample_idx"]), int(g["round"])],
                           "gap_index": j}
                    if index is not None:
                        row["bank_examples"] = bank_examples(index.retrieve(stmt, a.bank_k))
                    rows.append(row)
                sg_rows += rows
                plans.append({"parent_key": [g["problem_id"], int(g["sample_idx"]), int(g["round"])], "plan": SR.plan_to_dict(plan),
                              "sg_ids": [r["problem_id"] for r in rows]})
        st[route] += 1
        if route == "feedback":
            fb.append((g, v))
    _write(rd / "fb_parents.gen.jsonl", (g for g, _ in fb))
    _write(rd / "fb_parents.ver.jsonl", (v for _, v in fb))
    _write(rd / "sg_problems.jsonl", sg_rows)
    _write(rd / "sg_plans.jsonl", plans)
    (rd / "route.json").write_text(json.dumps(dict(st), indent=1))
    print(json.dumps(dict(st)))
    if probe_pool is not None:
        probe_pool.close()
    return 0


def _open_pool(workers: int):
    from align.verify.lean import ReplPool, find_toolchain
    from align.verify.lean.repl_pool import ReplLimits

    tc = find_toolchain()
    return ReplPool.from_toolchain(tc, n_workers=workers, limits=ReplLimits(), preflight=True,
                                   scratch_dir=str(Path(os.environ.get("TMPDIR", str(REPO / ".tmp"))) / "k5_route_pool"))


def _plan_with_pool(pool, statement: str, g: dict, v: dict, SR):
    """SubgoalPlan for a failed node, lemma statements checked to elaborate; None when Lean cannot run it."""
    from align.verify.lean.repl_pool import ReplError

    def work(w, _):
        try:
            probe = SR.repl_probe(w, statement)
            plan = SR.plan_subgoals(statement, g["lean_code"], v["repl_errors"], probe, parent_id=g["problem_id"], repl_code=v["repl_code"])
            for gap in plan.gaps:
                if gap.lemma is not None and not SR.check_lemma_elaborates(w, plan, gap):
                    gap.lemma, gap.why_not = None, "lemma_does_not_elaborate"
            return plan
        except (ReplError, ValueError):
            return None

    r = pool.map(work, [None])[0]
    return None if isinstance(r, BaseException) else r


def cmd_assemble(a) -> int:
    from align.rl.envs import subgoal_repair as SR

    rd = Path(a.run_dir, "agentic", f"round{a.round}")
    plans = _read(rd / "sg_plans.jsonl")
    sg_gens = _read(rd / "sg_gen.jsonl")
    sg_vers = _read(rd / "sg_ver.jsonl")
    ver_by = {C.gen_key(v): v for v in sg_vers}
    by_pid: dict[str, list[tuple[dict, dict]]] = defaultdict(list)
    for g in sg_gens:
        v = ver_by.get(C.gen_key(g))
        if v is None:
            raise ValueError(f"subgoal sample {C.gen_key(g)} is not verified: run the sg stage's verification to completion")
        by_pid[g["problem_id"]].append((g, v))
    # parents' generation records, from every pair of the previous round
    parents = {}
    for f in sorted(Path(a.run_dir, "agentic").glob("gen_r*.jsonl")) + sorted(Path(a.run_dir, "agentic").glob("round*/asm.gen.jsonl")):
        for g in _read(f):
            parents[(g["problem_id"], int(g["sample_idx"]), int(g["round"]))] = g
    pvers = {}
    for f in sorted(Path(a.run_dir, "agentic").glob("ver_r*.jsonl")) + sorted(Path(a.run_dir, "agentic").glob("round*/asm.ver.jsonl")):
        for v in _read(f):
            pvers[(v["problem_id"], int(v["sample_idx"]), int(v["round"]))] = v
    gens, vers, st = [], [], Counter()
    for pl in plans:
        pk = tuple(pl["parent_key"])
        pg, pv = parents[pk], pvers[pk]
        plan = SR.plan_from_dict(pl["plan"])
        bodies, sub_tokens = {}, 0
        for sid in pl["sg_ids"]:
            sub_tokens += sum(node_tokens(g) for g, _ in by_pid.get(sid, []))
        for j, gap in enumerate(plan.gaps):
            if gap.lemma is None:
                continue
            sid = sg_problem_id(pg, j)
            ok = [g for g, v in by_pid.get(sid, []) if v["accepted"] and g.get("lean_code")]
            ok.sort(key=lambda g: len(g["lean_code"]))
            for g in ok:
                body = SR.lemma_body(g["lean_code"], gap.lemma.name)
                if body is not None:
                    bodies[j] = body
                    break
        gen, ver = make_asm_nodes(plan, pg, pv, bodies, sub_tokens, len(bodies), a.round, a.cap)
        gens.append(gen)
        if ver is not None:
            vers.append(ver)
        st["assembled"] += 1
        st["gaps"] += len(plan.gaps)
        st["proven"] += len(bodies)
        st["unchanged_verdict_copied"] += ver is not None
    _write(rd / "asm.gen.jsonl", gens)
    if vers:
        _write(rd / "asm.ver.jsonl", vers)
    print(json.dumps(dict(st)))
    return 0


def _arm_files(run_dir: str, rounds: int) -> dict[str, tuple[list[str], list[str]]]:
    rd = Path(run_dir)
    ag = [rd / "agentic/gen_r0.jsonl"]
    av = [rd / "agentic/ver_r0.jsonl"]
    for r in range(1, rounds):
        ag += [rd / f"agentic/gen_r{r}.jsonl", rd / f"agentic/round{r}/asm.gen.jsonl"]
        av += [rd / f"agentic/ver_r{r}.jsonl", rd / f"agentic/round{r}/asm.ver.jsonl"]
    files = {
        "restart": ([rd / "restart/gen_r0.jsonl"], [rd / "restart/ver_r0.jsonl"]),
        "correct": ([rd / "restart/gen_r0.jsonl", rd / "correct/gen_r1.jsonl", rd / "correct/gen_r2.jsonl"],
                    [rd / "restart/ver_r0.jsonl", rd / "correct/ver_r1.jsonl", rd / "correct/ver_r2.jsonl"]),
        "agentic": (ag, av),
    }
    return {k: ([str(x) for x in g], [str(x) for x in v]) for k, (g, v) in files.items()}


def load_arms(run_dir: str, rounds: int = 4, arms: Iterable[str] = ARMS) -> dict[str, dict[str, list[Unit]]]:
    out = {}
    for arm, (gf, vf) in _arm_files(run_dir, rounds).items():
        if arm not in arms:
            continue
        if arm == "correct" and not Path(gf[1]).exists():
            continue  # no repair round generated yet: round-0 roots alone are not the correction arm
        gens = [g for f in gf for g in _read(f)]
        vers = [v for f in vf for v in _read(f)]
        if not gens:
            continue
        out[arm] = units_restart(gens, vers) if arm == "restart" else units_chains(gens, vers, arm)
    return out


def build_report(arms: dict[str, dict[str, list[Unit]]], budget_spec: str, n_perm: int = 256, seed: int = 0) -> dict:
    from align.eval.paired import holm, paired_bootstrap_ci, sign_flip_test

    if "restart" not in arms:
        raise ValueError("the restart arm defines the budget and the pairing; run it first")
    pids = sorted(arms["restart"])
    budgets = budget_table(budget_spec, arms["restart"], pids)
    rep = equalise(arms, budgets, n_perm, seed)
    out: dict = {"budget": budget_spec, "n_problems": len(pids), "n_perm": n_perm, "seed": seed, "arms": {}, "curves": {}}
    for arm in arms:
        out["arms"][arm] = arm_summary(rep[arm][1.0], budgets, 1.0)
        out["curves"][arm] = {str(f): arm_summary(rep[arm][f], budgets, f)["mean_p_solve"] for f in rep[arm]}
    fam = [("correct", "restart"), ("agentic", "restart"), ("agentic", "correct")]
    tests, ps = [], []
    for x, y in fam:
        if x in arms and y in arms:
            d = [rep[x][1.0][p].p_solve - rep[y][1.0][p].p_solve for p in pids]
            t = sign_flip_test(d, seed=seed)
            lo, hi = paired_bootstrap_ci(d, seed=seed)
            tests.append({"contrast": f"{x} - {y}", "mean_diff": t["observed"], "p": t["p"], "ci95": [lo, hi], "n": t["n"], "n_nonzero": t["n_nonzero"]})
            ps.append(t["p"])
    if ps:
        h = holm(ps)
        for t, adj, rej in zip(tests, h["adjusted"], h["reject"]):
            t["p_holm"], t["reject_at_0.05"] = adj, rej
    out["tests"] = tests
    out["family"] = "K5 pairwise (Holm over the contrasts present); the training-seed null of K0 does not apply: no training here"
    out["per_problem"] = {arm: {p: rep[arm][1.0][p].p_solve for p in pids} for arm in arms}
    return out


def cmd_report(a) -> int:
    arms = load_arms(a.run_dir, a.rounds)
    rep = build_report(arms, a.budget, a.n_perm, a.seed)
    Path(a.run_dir, "k5_report.json").write_text(json.dumps(rep, indent=1))
    for arm, s in rep["arms"].items():
        print(f"{arm:9s} p_solve@B={s['mean_p_solve']:.4f}  spent/B={s['spend_ratio']:.3f}  underfilled={s['underfilled_problems']}")
    for t in rep["tests"]:
        print(f"{t['contrast']:22s} diff={t['mean_diff']:+.4f} CI[{t['ci95'][0]:+.4f},{t['ci95'][1]:+.4f}] p_holm={t.get('p_holm', float('nan')):.4f}")
    return 0


def cmd_status(a) -> int:
    arms = load_arms(a.run_dir, a.rounds)
    pids = sorted(arms["restart"])
    budgets = budget_table(a.budget, arms["restart"], pids)
    for arm, by_p in arms.items():
        rep = {p: replay_budget(by_p.get(p, []), budgets[p], a.n_perm, a.seed, p) for p in pids}
        need = units_needed(by_p, budgets, rep)
        print(json.dumps({"arm": arm, "underfilled": len(need), "extra_units_total": sum(need.values()), "worst": sorted(need.items(), key=lambda kv: -kv[1])[:5]}))
    return 0


def cmd_ledger(a) -> int:
    arms = load_arms(a.run_dir, a.rounds)
    for arm, by_p in arms.items():
        rows = [dataclasses.asdict(u) for us in by_p.values() for u in us]
        _write(Path(a.run_dir, "ledger", f"units_{arm}.jsonl"), rows)
        print(arm, len(rows), "units over", len(by_p), "problems")
    return 0


def cmd_estimate(a) -> int:
    cfg = Cfg(run_dir="-", problems=[], n_restart=a.n_restart, n_agentic=a.n_agentic)
    sets = {"minif2f": a.minif2f, "proofnet": a.proofnet}
    print(json.dumps(estimate(sets, cfg), indent=1))
    return 0


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    def common(p, run=True):
        if run:
            p.add_argument("--run-dir", required=True)

    def cfgargs(p):
        d = Cfg(run_dir="", problems=[])
        common(p)
        p.add_argument("--problems", nargs="+", required=True)
        p.add_argument("--model", default=d.model)
        p.add_argument("--seed", type=int, default=d.seed)
        p.add_argument("--cap", type=int, default=d.cap)
        p.add_argument("--traj-cap", type=int, default=d.traj_cap)
        p.add_argument("--n-restart", type=int, default=d.n_restart)
        p.add_argument("--n-agentic", type=int, default=d.n_agentic)
        p.add_argument("--rounds", type=int, default=d.rounds)
        p.add_argument("--routes", default=d.routes)
        p.add_argument("--sub-n", type=int, default=d.sub_n)
        p.add_argument("--sub-cap", type=int, default=d.sub_cap)
        p.add_argument("--bank-k", type=int, default=d.bank_k)
        p.add_argument("--budget", default=d.budget)
        p.add_argument("--workers", type=int, default=d.workers)
        p.add_argument("--bank-problems")
        p.add_argument("--bank-gens")
        p.add_argument("--bank-vers")
        p.add_argument("--eval-problems", nargs="*", default=[])
        p.add_argument("--no-hostguard", dest="hostguard", action="store_false")

    p = sub.add_parser("plan", help="print the stages and exact command lines")
    cfgargs(p)
    p = sub.add_parser("run", help="execute the stages (resumable)")
    cfgargs(p)
    p.add_argument("--only", nargs="*", default=None)

    p = sub.add_parser("prep")
    common(p)
    p.add_argument("--problems", nargs="+", required=True)
    p.add_argument("--bank-problems")
    p.add_argument("--bank-gens")
    p.add_argument("--bank-vers")
    p.add_argument("--eval-problems", nargs="*", default=[])
    p.add_argument("--bank-k", type=int, default=3)

    p = sub.add_parser("route")
    common(p)
    p.add_argument("--round", type=int, required=True)
    p.add_argument("--mode", choices=("subgoal", "feedback"), required=True)
    p.add_argument("--pairs", nargs=2, action="append", metavar=("VER", "GEN"), required=True)
    p.add_argument("--workers", type=int, default=2)
    p.add_argument("--bank-k", type=int, default=3)
    p.add_argument("--sub-n", type=int, default=4)

    p = sub.add_parser("assemble")
    common(p)
    p.add_argument("--round", type=int, required=True)
    p.add_argument("--cap", type=int, default=40960)

    for name, fn in (("report", cmd_report), ("status", cmd_status), ("ledger", cmd_ledger)):
        p = sub.add_parser(name)
        common(p)
        p.add_argument("--budget", default="restart_k:8")
        p.add_argument("--rounds", type=int, default=4)
        p.add_argument("--n-perm", type=int, default=256)
        p.add_argument("--seed", type=int, default=0)

    p = sub.add_parser("estimate")
    p.add_argument("--minif2f", type=int, default=226)
    p.add_argument("--proofnet", type=int, default=185)
    p.add_argument("--n-restart", type=int, default=32)
    p.add_argument("--n-agentic", type=int, default=8)

    a = ap.parse_args(argv)
    if a.cmd in ("plan", "run"):
        cfg = Cfg(run_dir=a.run_dir, problems=a.problems, model=a.model, seed=a.seed, cap=a.cap, traj_cap=a.traj_cap, n_restart=a.n_restart,
                  n_agentic=a.n_agentic, rounds=a.rounds, routes=a.routes, sub_n=a.sub_n, sub_cap=a.sub_cap, bank_k=a.bank_k, budget=a.budget,
                  workers=a.workers, bank_problems=a.bank_problems, bank_gens=a.bank_gens, bank_vers=a.bank_vers,
                  eval_problems=a.eval_problems, hostguard=a.hostguard)
        stages = build_stages(cfg)
        if a.cmd == "plan":
            for st in stages:
                print(f"# {st.name} [{'GPU' if st.gpu else 'CPU'}]\n{' '.join(shlex.quote(x) for x in _hg(cfg, st))}\n")
            return 0
        return run_stages(cfg, stages, Runner(), only=set(a.only) if a.only else None)
    return {"prep": cmd_prep, "route": cmd_route, "assemble": cmd_assemble, "report": cmd_report, "status": cmd_status,
            "ledger": cmd_ledger, "estimate": cmd_estimate}[a.cmd](a)


if __name__ == "__main__":
    raise SystemExit(main())