onf-c0/src/onf/graph/cli.py

"""python -m onf.graph — the one entrypoint for the demonstration-graph subsystem's offline
pipeline: build (hdf5 -> g_nodes/g_edges), calibrate (the corpus's own constants, see onf.calib)
and train (the retrieval-network curriculum, see onf.graph.train.loop). This is the repo's only
supported CLI surface for this subsystem.

WHY A CLI, NOT A PILE OF SCRIPTS
    Both library functions this module calls already do the real work (onf.graph.build.from_demos.
    build_graph, onf.graph.train.loop.train_graph) and already accept a onf.graph.report.
    StageLogger. What was missing was a single, reproducible, un-hardcoded front door: every subcommand
    here resolves its FULL configuration (suite -> hdf5 dir -> GraphConfig, never a literal path),
    can be inspected with --dry-run before anything expensive runs, and always leaves a
    config/manifest/metrics/log record under outputs/ -- the project's own reproducibility bar.

WHY EVERY SUBCOMMAND OWNS A StageLogger, BUT ONLY build REPOINTS latest
    onf.config.Paths.graph resolves to outputs/<suite>/latest/artifacts -- "the graph a
    suite's train run should read". A graph BUILD is the event that pointer exists to track, so
    build's StageLogger repoints it (the default). train writes its checkpoint straight
    into that SAME resolved directory (mirroring onf.graph.train.loop.train_graph's own contract: the
    checkpoint lives next to the g_nodes.npz/g_edges.npz it was fitted against, not off in its
    own run dir) and passes update_latest=False to its own bookkeeping StageLogger -- a
    train run would otherwise silently orphan the graph it just extended from Paths.graph's point
    of view.

SHAPE
    SuiteCommand owns the whole per-suite lifecycle every subcommand shares -- resolve config,
    --dry-run echo, precondition check, --force/skip guard, StageLogger, summary table -- and
    BuildCommand / TrainCommand supply only what actually differs between them.
    Adding a subcommand means one subclass plus one add_parser block, not a fourth copy of that
    lifecycle.

EXIT CODES (every subcommand)
    0 success -- 1 an exception / bad input -- 2 a pre-registered gate failed, which only
    calibrate has (build and train have none).
"""

from __future__ import annotations

import argparse
import dataclasses
import glob
import json
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, ClassVar

import numpy as np
import yaml

from onf.config import GraphConfig, Paths, default_paths
from onf.graph.core import schema
from onf.graph.build.from_demos import LoadLimits, _hdf5_name, build_graph
from onf.graph.report import StageLogger

# ---- repo-relative paths (the ONE place this module resolves a filesystem literal; everything else
# goes through onf.config.Paths / configs/suites.yaml, per the module docstring). Mirrors onf.config's
# own "parents[N] == repo root" idiom. ---------------------------------------------------------------
_REPO_ROOT = Path(__file__).resolve().parents[3]        # src/onf/graph/cli.py -> repo root
_SUITES_YAML = _REPO_ROOT / "configs" / "suites.yaml"

# build's internal pipeline steps, in order (matches onf.graph.build.from_demos.build_graph /
# onf.graph.core.edges.EdgeSet.build's own step calls -- no boundary step any more, since there is no
# stage_next/stage_prev relation to build; see schema.RELATIONS's module note). Used both to
# print "planned stages" under --dry-run and as this CLI's own StageLogger's n_stages; the real
# run reports the same labels itself via StageLogger.step().
BUILD_STAGE_LABELS: tuple[str, ...] = (
    "load/hdf5", "nodes/coarsen", "temporal", "sibling", "align", "edges/summary",
)
DEFAULT_TRAIN_EPOCHS = 30   # a full curriculum epoch count for one stage (train.py's own runs use this)
# The only legal stages value -- see onf.graph.train.loop.train_graph.
TRAIN_STAGES: tuple[int, ...] = (1,)


class CLIError(Exception):
    """Raised by any subcommand to fail cleanly with a specific exit code (1 by default) instead of a
    raw traceback -- main is the only place this is caught."""

    def __init__(self, message: str, exit_code: int = 1):
        super().__init__(message)
        self.exit_code = exit_code


# ======================================================================================================
# suites.yaml -- the single source of "which suites this CLI knows about" (never a literal list)
# ======================================================================================================
def load_configured_suites(path: str | os.PathLike | None = None) -> list[str]:
    """The suite names --suite all expands to.

    Args:
        path: Override for configs/suites.yaml.

    Returns:
        The suites: mapping keys, in file order. Never hardcoded here -- add a suite to that YAML
        and this CLI (and every downstream launcher that already reads it) picks it up unchanged.

    Raises:
        CLIError: The file is missing or has no suites: mapping.
    """
    p = Path(path) if path is not None else _SUITES_YAML
    if not p.exists():
        raise CLIError(f"suites config not found: {p}")
    with open(p) as f:
        doc = yaml.safe_load(f) or {}
    suites = doc.get("suites") or {}
    if not suites:
        raise CLIError(f"{p}: no `suites:` mapping found")
    return list(suites.keys())


def _resolve_suites(suite_arg: str, configured: list[str]) -> list[str]:
    """--suite -> the list of suites to act on. "all" expands to every configured suite;
    anything else must be a name FROM that list -- an unrecognised suite fails loudly here rather than
    reaching onf.config.Paths.hdf5 and failing with a confusing "no such directory"."""
    if suite_arg == "all":
        return list(configured)
    if suite_arg not in configured:
        raise CLIError(
            f"unknown suite {suite_arg!r}; choices: {sorted(configured)} or 'all' "
            f"(from {_SUITES_YAML})"
        )
    return [suite_arg]


# ======================================================================================================
# small shared parsing / formatting helpers
# ======================================================================================================
def _json_default(obj: Any) -> Any:
    """json.dumps fallback: numpy scalars/arrays and Path -> plain JSON, anything else -> str."""
    if isinstance(obj, np.generic):
        return obj.item()
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, Path):
        return os.fspath(obj)
    return str(obj)


def _echo_config(label: str, config: dict) -> None:
    """Print one fully-resolved config block to stdout."""
    print(f"===== {label}: fully-resolved config =====")
    print(json.dumps(config, indent=2, default=_json_default, sort_keys=True))


def _print_table(headers: list[str], rows: list[list[str]]) -> None:
    """A minimal aligned text table -- no external dependency, used by build/train. Column
    widths are the max of the header and every row's cell in that column."""
    if not rows:
        print("  (no rows)")
        return
    widths = [len(h) for h in headers]
    for row in rows:
        for i, cell in enumerate(row):
            widths[i] = max(widths[i], len(str(cell)))
    line = "  ".join(h.ljust(widths[i]) for i, h in enumerate(headers))
    print(line)
    print("  ".join("-" * w for w in widths))
    for row in rows:
        print("  ".join(str(cell).ljust(widths[i]) for i, cell in enumerate(row)))


def _fmt_num(x: Any, nd: int = 4) -> str:
    """One scalar -> a fixed-precision display string; non-finite floats print verbatim."""
    if isinstance(x, bool):
        return str(x)
    if isinstance(x, float):
        if not np.isfinite(x):
            return str(x)
        return f"{x:.{nd}f}"
    return str(x)


def _fmt_metric(v: Any) -> str:
    """A single metrics.json leaf -> a compact display string. A bootstrap CI dict
    ({"lo":...,"point":...,"hi":...}) prints as point [lo, hi] -- never just the point estimate
    alone, so a report never hides the uncertainty the probe itself computed."""
    if isinstance(v, dict) and {"lo", "point", "hi"} <= set(v.keys()):
        return f"{_fmt_num(v['point'])} [{_fmt_num(v['lo'])}, {_fmt_num(v['hi'])}]"
    if isinstance(v, dict):
        return "{...}"     # nested, non-CI dict -- not flattened at this level
    if isinstance(v, list):
        return f"[{len(v)} items]" if len(v) > 6 else str(v)
    return _fmt_num(v)


def _read_json(path: Path) -> dict:
    """Read a JSON object, degrading to {} when the file is missing, truncated, or unreadable --
    a half-written run record must never crash the scan that walks past it."""
    if not path.exists():
        return {}
    try:
        loaded = json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        return {}
    return loaded if isinstance(loaded, dict) else {}


# ======================================================================================================
# outputs/ run-directory scanning (the generic "skip if already done" check every subcommand uses)
# ======================================================================================================
def _find_matching_run(root: Path, suite: str, stage: str, config: dict) -> Path | None:
    """The generic "already did this exact work" check every subcommand's --force guards.

    Args:
        root: Run-bookkeeping root (outputs/).
        suite: Suite short-name.
        stage: StageLogger stage label (graph / train).
        config: The fully-resolved config to match against.

    Returns:
        The newest COMPLETE run under <root>/<suite>/<stage>_* whose config.json matches
        config -- ignoring the purely cosmetic name key, since two runs with the same knobs and
        different --name labels are still the same work -- or None.
    """
    base = root / suite
    if not base.is_dir():
        return None
    norm_target = {k: v for k, v in config.items() if k != "name"}
    for run_dir in sorted(base.glob(f"{stage}_*"), reverse=True):
        if not run_dir.is_dir() or run_dir.is_symlink():
            continue
        if _read_json(run_dir / "manifest.json").get("status") != "complete":
            continue
        norm_got = {k: v for k, v in _read_json(run_dir / "config.json").items() if k != "name"}
        if norm_got == norm_target:
            return run_dir
    return None


# ======================================================================================================
# --resume: inherit a prior run's resolved config for any flag the user did not explicitly pass
# ======================================================================================================
def _load_resume_config(run_dir: str | os.PathLike) -> dict:
    """Load a prior StageLogger run's config.json for --resume layering.

    Raises:
        CLIError: run_dir has no config.json, or it is unreadable/empty.
    """
    p = Path(run_dir) / "config.json"
    if not p.exists():
        raise CLIError(f"--resume {run_dir}: no config.json found there (not a StageLogger run dir?)")
    config = _read_json(p)
    if not config:
        raise CLIError(f"--resume {run_dir}: unreadable or empty config.json")
    return config


def _pick(explicit: Any, resumed: dict, key: str, default: Any) -> Any:
    """explicit (a CLI value, or argparse's None sentinel meaning "not passed") wins; else the
    --resumed run's own resolved value for key; else default. The layering that makes
    --resume meaningful: a caller can reproduce a prior run exactly, or reproduce it while
    overriding exactly the flags they pass this time."""
    if explicit is not None:
        return explicit
    if key in resumed:
        return resumed[key]
    return default


# ======================================================================================================
# per-suite outcome
# ======================================================================================================
@dataclass
class SuiteResult:
    """What one subcommand did for one suite.

    Attributes:
        suite: Suite short-name.
        status: dry-run | skip | the subclass's done_status | failed.
        run_dir: The matching run directory, when this was a skip.
        metrics: Whatever SuiteCommand.execute returned; drives the summary row.
    """

    suite: str
    status: str
    run_dir: Path | None = None
    metrics: dict[str, Any] = field(default_factory=dict)

    @property
    def ok(self) -> bool:
        """Whether this suite counts toward a zero exit code."""
        return self.status != "failed"


# ======================================================================================================
# subcommands
# ======================================================================================================
class SuiteCommand:
    """The per-suite lifecycle every subcommand shares, owning all resolved CLI state.

    One instance per main() invocation. run walks the selected suites; run_suite
    is the fixed pipeline -- resolve config -> --dry-run echo -> validate ->
    --force/skip guard -> StageLogger -> execute. Subclasses override only the hooks.
    """

    # StageLogger stage label, and the <stage>_* run-dir prefix the skip guard scans.
    stage: ClassVar[str] = ""
    # Human-facing subcommand name used in progress/error lines.
    label: ClassVar[str] = ""
    # The summary-table status for a suite that actually ran.
    done_status: ClassVar[str] = "done"
    # Whether this stage's StageLogger repoints <suite>/latest (see the module docstring).
    update_latest: ClassVar[bool] = True
    # Summary-table header row; row must return this many cells.
    table_headers: ClassVar[tuple[str, ...]] = ()

    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.paths: Paths = default_paths()
        self._resumed = _load_resume_config(args.resume) if getattr(args, "resume", None) else {}
        self.suites = _resolve_suites(args.suite, load_configured_suites())
        # Paths.outputs() only RESOLVES a path (it never mkdirs), so computing this eagerly is still
        # compatible with --dry-run's "touches NOTHING under outputs/" contract.
        self.out_root = Path(args.out) if args.out else self.paths.outputs()

    # ---- resolution ------------------------------------------------------------------------------
    def opt(self, key: str, default: Any) -> Any:
        """One knob, resolved: explicit CLI flag > --resumed value > default (see _pick)."""
        return _pick(getattr(self.args, key, None), self._resumed, key, default)

    # ---- driver ----------------------------------------------------------------------------------
    def run(self) -> int:
        """Run every selected suite, then print the summary table.

        Returns:
            The process exit code: 0 when every suite succeeded, else 1.

        Raises:
            CLIError: A precondition failed (bad input / missing artifact). Unlike an unexpected
                exception -- which is reported per-suite and lets the remaining suites continue -- a
                CLIError is a config-level problem, so it aborts the whole invocation for main to
                turn into its own exit code.
        """
        results: list[SuiteResult] = []
        for suite in self.suites:
            try:
                results.append(self.run_suite(suite))
            except CLIError:
                raise
            except Exception as e:
                print(f"[{self.label}] {suite}: FAILED -- {type(e).__name__}: {e}", file=sys.stderr)
                results.append(SuiteResult(suite, "failed"))

        if self.args.dry_run:
            return 0
        _print_table(list(self.table_headers), [self.row(r) for r in results])
        return 0 if all(r.ok for r in results) else 1

    def run_suite(self, suite: str) -> SuiteResult:
        """Resolve, guard, and execute this subcommand for one suite."""
        config = self.resolve_config(suite)

        if self.args.dry_run:
            _echo_config(f"{self.label}[{suite}]", config)
            print(self.plan_line(suite, config))
            return SuiteResult(suite, "dry-run")

        self.validate(suite, config)

        if not self.args.force:
            hit = _find_matching_run(self.out_root, suite, self.stage, config)
            if hit is not None and self.artifacts_present(suite, config):
                print(f"[{self.label}] {suite}: SKIP -- {self.skip_note(hit, config)} "
                      f"(pass --force to re-run)")
                return SuiteResult(suite, "skip", run_dir=hit)

        _echo_config(f"{self.label}[{suite}]", config)
        with self._logger(suite, config) as log:
            metrics = self.execute(suite, config, log)
        return SuiteResult(suite, self.done_status, metrics=metrics)

    def _logger(self, suite: str, config: dict) -> StageLogger:
        """This run's StageLogger, built from the shared flags plus the per-command hooks."""
        return StageLogger(
            self.stage, suite, name=self.args.name, root=self.args.out, config=config,
            n_stages=self.n_stages(config), inputs=self.stage_inputs(suite, config),
            quiet=self.args.quiet, update_latest=self.update_latest,
        )

    # ---- hooks -----------------------------------------------------------------------------------
    def resolve_config(self, suite: str) -> dict:
        """The fully-resolved config for one suite -- every default materialized, no bare placeholders.
        This dict is what lands in config.json and what the skip guard compares against."""
        raise NotImplementedError

    def execute(self, suite: str, config: dict, log: StageLogger) -> dict:
        """Do the actual work inside an entered StageLogger. Returns the summary-row metrics."""
        raise NotImplementedError

    def row(self, result: SuiteResult) -> list[str]:
        """One SuiteResult -> one summary-table row matching table_headers."""
        raise NotImplementedError

    def plan_line(self, suite: str, config: dict) -> str:
        """The --dry-run "planned stages" line printed under the resolved config."""
        return f"planned stages ({self.n_stages(config)})"

    def n_stages(self, config: dict) -> int:
        """How many StageLogger.step calls this run is expected to make."""
        return 0

    def stage_inputs(self, suite: str, config: dict) -> list[str]:
        """Input artifact paths to fingerprint into the run manifest."""
        return []

    def validate(self, suite: str, config: dict) -> None:
        """Fail loudly, before any expensive work, when a precondition is unmet."""

    def artifacts_present(self, suite: str, config: dict) -> bool:
        """Whether the on-disk artifacts a matching prior run should have left still exist. False
        demotes a config match back to a real run -- a matching record whose output was deleted is not
        work that is actually done."""
        return True

    def skip_note(self, hit: Path, config: dict) -> str:
        """The middle of the SKIP line: what was found, and where."""
        return f"identical config already run at {hit}"


class BuildCommand(SuiteCommand):
    """HDF5 demos -> g_nodes.npz / g_edges.npz for one suite."""

    stage = "graph"
    label = "build"
    done_status = "built"
    update_latest = True                # a build IS the "new graph" event latest exists to track
    table_headers = ("suite", "status", "n_nodes", "n_edges", "n_stages")

    def resolve_config(self, suite: str) -> dict:
        cfg = GraphConfig(
            coarsen=self.opt("coarsen", schema.COARSEN),
            device=self.opt("device", ""),
        )
        hdf5_dir = self.opt("hdf5_dir", None)
        hdf5_dir = os.fspath(hdf5_dir) if hdf5_dir else os.fspath(self.paths.hdf5(_hdf5_name(suite)))
        limits = {
            k: v for k, v in (
                ("limit_demos", self.opt("limit_demos", None)),
                ("limit_tasks", self.opt("limit_tasks", None)),
            ) if v is not None
        }
        # GraphConfig is splatted FIRST so the explicitly-resolved keys below always win. This is
        # load-bearing, not cosmetic: GraphConfig carries its own graph_dir (onf.config's
        # GR_GRAPH_DIR override) and TrainCommand resolves a key of the same name -- splatting last
        # would overwrite the resolved path with GraphConfig's default "".
        return {
            **dataclasses.asdict(cfg),
            "command": "build", "suite": suite, "hdf5_dir": hdf5_dir,
            "seed": self.opt("seed", 0), "name": self.args.name, "limits": limits,
        }

    def n_stages(self, config: dict) -> int:
        return len(BUILD_STAGE_LABELS)

    def plan_line(self, suite: str, config: dict) -> str:
        return f"planned stages ({len(BUILD_STAGE_LABELS)}): {', '.join(BUILD_STAGE_LABELS)}"

    def stage_inputs(self, suite: str, config: dict) -> list[str]:
        return sorted(glob.glob(f"{config['hdf5_dir']}/*.hdf5"))

    def skip_note(self, hit: Path, config: dict) -> str:
        return f"identical config already built at {hit}"

    def execute(self, suite: str, config: dict, log: StageLogger) -> dict:
        cfg = GraphConfig(coarsen=config["coarsen"], device=config["device"])
        limits = config["limits"]
        artifacts = build_graph(
            suite, hdf5_dir=config["hdf5_dir"], out_dir=None, cfg=cfg, coarsen=cfg.coarsen,
            device=(cfg.device or None), logger=log,
            limits=LoadLimits(tasks=limits.get("limit_tasks"), demos=limits.get("limit_demos")),
        )
        return artifacts.to_dict()

    def row(self, result: SuiteResult) -> list[str]:
        m = result.metrics
        return [result.suite, result.status,
                str(m.get("n_nodes", "-")), str(m.get("n_edges", "-")), str(m.get("n_stages", "-"))]


class TrainCommand(SuiteCommand):
    """Train GraphRetrieverNet on a built graph -> g_head.npz next to that graph."""

    stage = "train"
    label = "train"
    done_status = "trained"
    update_latest = False               # a train run must not steal latest from the build it read
    table_headers = ("suite", "status", "top1", "mrr", "crossed_top1")

    def resolve_config(self, suite: str) -> dict:
        # Imported here, not at module scope: python -m onf.graph --help and every build run would
        # otherwise pay for importing torch through onf.graph.train.loop.
        from onf.graph.train.loop import N_EVAL, N_QUERY
        from onf.graph.train.loss import RETRIEVAL

        cfg = GraphConfig(device=self.opt("device", ""))
        graph_dir = self.opt("graph_dir", None)
        graph_dir = Path(graph_dir) if graph_dir else Path(self.paths.graph(suite))
        # asdict FIRST -- see BuildCommand.resolve_config: GraphConfig.graph_dir would otherwise
        # overwrite the resolved graph_dir below with "".
        return {
            **dataclasses.asdict(cfg),
            "command": "train", "suite": suite, "graph_dir": str(graph_dir),
            "stages": list(TRAIN_STAGES), "epochs": self.opt("epochs", DEFAULT_TRAIN_EPOCHS),
            "n_query": self.opt("n_query", N_QUERY), "n_eval": self.opt("n_eval", N_EVAL),
            "objective": self.opt("objective", RETRIEVAL),
            "seed": self.opt("seed", 0), "name": self.args.name,
        }

    def n_stages(self, config: dict) -> int:
        return len(config["stages"])

    def plan_line(self, suite: str, config: dict) -> str:
        return f"planned stages: {tuple(config['stages'])} (graph_dir={config['graph_dir']})"

    def stage_inputs(self, suite: str, config: dict) -> list[str]:
        return [str(Path(config["graph_dir"]) / schema.EDGES_NPZ)]

    def validate(self, suite: str, config: dict) -> None:
        from onf.graph.train.loss import CHUNK, RETRIEVAL

        graph_dir = Path(config["graph_dir"])
        if not (graph_dir / schema.NODES_NPZ).exists():
            raise CLIError(f"no built graph for suite {suite!r} at {graph_dir} -- run "
                            f"`python -m onf.graph build --suite {suite}` first")
        # Checked here rather than left to argparse choices, which would need the two constants
        # spelled as literals at parser-build time just to avoid importing torch.
        if config["objective"] not in (RETRIEVAL, CHUNK):
            raise CLIError(f"unknown --objective {config['objective']!r}; "
                            f"choices: {RETRIEVAL!r}, {CHUNK!r}")

    def artifacts_present(self, suite: str, config: dict) -> bool:
        return (Path(config["graph_dir"]) / schema.HEAD_NPZ).exists()

    def skip_note(self, hit: Path, config: dict) -> str:
        checkpoint = Path(config["graph_dir"]) / schema.HEAD_NPZ
        return f"identical config already trained at {hit}, checkpoint={checkpoint}"

    def execute(self, suite: str, config: dict, log: StageLogger) -> dict:
        from onf.graph.train.loop import train_graph
        from onf.graph.train.types import TrainGraphSpec

        cfg = GraphConfig(device=config["device"])
        result = train_graph(
            Path(config["graph_dir"]),
            TrainGraphSpec(stages=tuple(config["stages"]), epochs=config["epochs"],
                            seed=config["seed"], device=(config["device"] or None), save=True,
                            cfg=cfg, logger=log, n_query=config["n_query"],
                            n_eval=config["n_eval"], suite=suite,
                            objective=config["objective"]),
        )
        log.metric(checkpoint_path=result["checkpoint_path"], graph_hash=result["graph_hash"])
        return {"checkpoint_path": result["checkpoint_path"], **self._final_epoch(result["history"])}

    @staticmethod
    def _final_epoch(history: dict[Any, list[Any]]) -> dict:
        """The last stage's epoch == "final" record out of a training history, or {}."""
        final: dict = {}
        for stage_hist in history.values():
            if stage_hist and isinstance(stage_hist[-1], dict) and stage_hist[-1].get("epoch") == "final":
                final = stage_hist[-1]
        return final

    def row(self, result: SuiteResult) -> list[str]:
        m = result.metrics
        return [result.suite, result.status,
                _fmt_metric(m.get("top1", "-")), _fmt_metric(m.get("mrr", "-")),
                _fmt_metric(m.get("crossed_top1", "-"))]


class CalibrateCommand(SuiteCommand):
    """C0: demo corpus -> constants.json + action_scale.json beside the graph (onf.calib)."""

    stage = "calibrate"
    label = "calibrate"
    done_status = "calibrated"
    update_latest = False               # C0 reads the graph latest points at; it does not build one
    table_headers = ("suite", "status", "pos_scale", "kernel_bw", "assumed")

    def resolve_config(self, suite: str) -> dict:
        from onf.calib.measure import CalibrationSpec

        graph_dir = self.opt("graph_dir", None)
        graph_dir = Path(graph_dir) if graph_dir else Path(self.paths.graph(suite))
        hdf5_dir = self.opt("hdf5_dir", None)
        hdf5_dir = os.fspath(hdf5_dir) if hdf5_dir else os.fspath(self.paths.hdf5(_hdf5_name(suite)))
        return {
            "command": "calibrate", "suite": suite, "graph_dir": str(graph_dir),
            "hdf5_dir": hdf5_dir, "name": self.args.name,
            "spec": dataclasses.asdict(CalibrationSpec(n_queries=self.opt("n_queries", 256))),
        }

    def n_stages(self, config: dict) -> int:
        return 1

    def plan_line(self, suite: str, config: dict) -> str:
        return f"planned stages (1): measure (graph_dir={config['graph_dir']})"

    def stage_inputs(self, suite: str, config: dict) -> list[str]:
        return [str(Path(config["graph_dir"]) / schema.NODES_NPZ)]

    def validate(self, suite: str, config: dict) -> None:
        if not (Path(config["graph_dir"]) / schema.NODES_NPZ).exists():
            raise CLIError(f"no built graph for suite {suite!r} at {config['graph_dir']} -- run "
                            f"`python -m onf.graph build --suite {suite}` first")

    def artifacts_present(self, suite: str, config: dict) -> bool:
        from onf.calib.constants import CONSTANTS_JSON

        return (Path(config["graph_dir"]) / CONSTANTS_JSON).exists()

    def execute(self, suite: str, config: dict, log: StageLogger) -> dict:
        from onf.calib.calibrator import Calibrator
        from onf.calib.constants import ASSUMED
        from onf.calib.corpus import DemoCorpus
        from onf.calib.measure import CalibrationSpec

        spec = CalibrationSpec(**config["spec"])
        corpus = DemoCorpus(suite, hdf5_dir=config["hdf5_dir"], graph_dir=config["graph_dir"],
                            paths=self.paths)
        log.step("measure")
        result = Calibrator(corpus, spec).run()
        constants = result.constants
        log.metric(content_hash=constants.content_hash,
                    constants_path=str(result.constants_path),
                    action_scale_path=str(result.action_scale_path))
        for name in sorted(constants.entries):
            entry = constants[name]
            log.log(f"{name} = {entry.value!r} [{entry.kind}]")
        for failure in result.gate_failures:
            log.log(f"GATE FAILED -- {failure}")
        # Written first, refused second: the numbers are the corpus's answer either way, and a run
        # that cannot be inspected cannot be argued with.
        if not result.ok:
            raise CLIError(
                f"{suite}: {len(result.gate_failures)} pre-registered gate(s) failed after writing "
                f"{result.constants_path}:\n  " + "\n  ".join(result.gate_failures),
                exit_code=2,
            )
        return {
            "pos_scale": constants.value("pos_scale"), "kernel_bw": constants.value("kernel_bw"),
            "assumed": ", ".join(constants.of_kind(ASSUMED)) or "-",
            "content_hash": constants.content_hash,
        }

    def row(self, result: SuiteResult) -> list[str]:
        m = result.metrics
        return [result.suite, result.status,
                _fmt_metric(m.get("pos_scale", "-")), _fmt_metric(m.get("kernel_bw", "-")),
                str(m.get("assumed", "-"))]


# ======================================================================================================
# top-level parser / dispatch
# ======================================================================================================
def _add_common_args(sub: argparse.ArgumentParser) -> None:
    """Attach the run-bookkeeping flags every subcommand shares."""
    sub.add_argument("--out", default=None, help="root run-bookkeeping dir (default: Paths.outputs(), "
                      "i.e. $ONF_OUTPUTS or <repo>/outputs)")
    sub.add_argument("--name", default="", help="label appended to this run's directory name")
    sub.add_argument("--seed", type=int, default=None, help="RNG seed (default: 0, or the resumed "
                      "run's seed under --resume)")
    sub.add_argument("--device", default=None, help='torch device, e.g. "cuda:0" or "cpu" '
                      '(default: auto -- cuda if available else cpu)')
    sub.add_argument("--quiet", action="store_true", help="suppress the mirrored stdout progress "
                      "lines (log.txt under outputs/ still gets every line)")
    sub.add_argument("--dry-run", action="store_true", help="print the fully-resolved config and the "
                      "planned stages, then exit 0 -- touches NOTHING under outputs/, so an expensive "
                      "run can be inspected before it starts")
    sub.add_argument("--force", action="store_true", help="re-run even if a completed run with the "
                      "identical resolved config already exists under outputs/")
    sub.add_argument("--resume", default=None, metavar="RUN_DIR", help="inherit RUN_DIR's "
                      "resolved config.json as defaults for any flag not explicitly passed here")


def build_parser() -> argparse.ArgumentParser:
    """Build the top-level parser.

    Returns:
        A parser whose namespace carries command_cls -- the SuiteCommand subclass
        main instantiates and runs.
    """
    p = argparse.ArgumentParser(
        prog="python -m onf.graph",
        description="The demonstration-graph subsystem's offline pipeline (see onf.graph.core.schema for "
        "the full WHY): build the graph, then train the retrieval network on it. Every subcommand "
        "resolves its full configuration through onf.config.Paths / configs/suites.yaml (never a "
        "hardcoded path) and leaves a reproducible run under outputs/ (onf.graph.report.StageLogger).",
    )
    sub = p.add_subparsers(dest="command", required=True, metavar="command")

    build = sub.add_parser(
        "build", help="HDF5 demos -> g_nodes.npz / g_edges.npz",
        description="Build the demonstration graph for one suite (or --suite all): coarsen raw demo "
        "frames into nodes (splitting only at gripper-state changes), then build the R=12-relation "
        "edge set (onf.graph.core.schema). `align` is different task, same phase bin, same gripper state "
        "-- there is deliberately no `stage` constraint here (see schema.RELATIONS's module note); a "
        "bad build here silently degrades `align`/`sibling` for every downstream stage.",
    )
    build.add_argument("--suite", required=True, help="object | spatial | goal | long | all "
                        "(configs/suites.yaml)")
    build.add_argument("--coarsen", type=int, default=None, help=f"raw frames per node "
                        f"(default: schema.COARSEN={schema.COARSEN})")
    build.add_argument("--hdf5-dir", default=None, help="override the demo hdf5 directory (default: "
                        "Paths.hdf5(<suite>) -- pass this to build from a fixture/test dataset)")
    build.add_argument("--limit-demos", type=int, default=None, help="cap demos read per hdf5 file "
                        "(smoke-test a suite without the full 500 demos)")
    build.add_argument("--limit-tasks", type=int, default=None, help="cap hdf5 files (tasks) read")
    _add_common_args(build)
    build.set_defaults(command_cls=BuildCommand)

    calibrate = sub.add_parser(
        "calibrate", help="C0: demo corpus -> constants.json (onf.calib)",
        description="Measure this suite's constants off its own corpus and write constants.json "
        "(schema-versioned, content-hashed) plus action_scale.json beside the graph artifacts. No "
        "parameters and no gradient: every number is the output of one measuring function, so a new "
        "suite is a calibration run rather than a code edit. Exit 2 means a pre-registered gate "
        "failed (action-scale R2, or a dilation schedule too shallow to cross the median demo).",
    )
    calibrate.add_argument("--suite", required=True, help="object | spatial | goal | long | all")
    calibrate.add_argument("--graph-dir", default=None, help="override the graph dir to calibrate "
                            "against and write into (default: Paths.graph(<suite>))")
    calibrate.add_argument("--hdf5-dir", default=None, help="override the demo hdf5 directory "
                            "(default: Paths.hdf5(<suite>))")
    calibrate.add_argument("--n-queries", type=int, default=None, help="calibration queries for the "
                            "seed/posterior support measurement (default: 256)")
    _add_common_args(calibrate)
    calibrate.set_defaults(command_cls=CalibrateCommand)

    train = sub.add_parser(
        "train", help="train GraphRetrieverNet (onf.graph.train.loop) on a built graph",
        description="Run the demo-graph retrieval curriculum (stage 1, temporal-advance pretraining, "
        "the only stage) and write g_head.npz next to the graph it was trained on, or with "
        "--objective chunk train the blend weight against a frozen retriever into g_alpha.npz. "
        "Requires `build` to have run first.",
    )
    train.add_argument("--suite", required=True, help="object | spatial | goal | long | all")
    train.add_argument("--epochs", type=int, default=None, help=f"epochs per stage "
                        f"(default: {DEFAULT_TRAIN_EPOCHS})")
    train.add_argument("--n-query", type=int, default=None, help="training queries per stage "
                        "(default: onf.graph.train.N_QUERY)")
    train.add_argument("--n-eval", type=int, default=None, help="held-out queries for per-epoch metrics "
                        "(default: onf.graph.train.N_EVAL)")
    train.add_argument("--graph-dir", default=None, help="override the graph dir to train on/into "
                        "(default: Paths.graph(<suite>))")
    train.add_argument("--objective", default=None, help="which loss to train: retrieval (stage 1's "
                        "WHEN/WHERE factorisation, the default) or chunk (stage 2's "
                        "mse(a_exec, a_corr) -- freezes the retriever, trains the blend weight and "
                        "writes g_alpha.npz; needs a_pi_raw.npz and action_scale.json beside the "
                        "graph)")
    _add_common_args(train)
    train.set_defaults(command_cls=TrainCommand)

    return p


def main(argv: list[str] | None = None) -> int:
    """CLI entrypoint.

    Args:
        argv: Argument list; None uses sys.argv[1:].

    Returns:
        The process exit code (0 success, 1 failure).
    """
    args = build_parser().parse_args(sys.argv[1:] if argv is None else argv)
    try:
        return args.command_cls(args).run()
    except CLIError as e:
        print(f"error: {e}", file=sys.stderr)
        return e.exit_code
    except FileNotFoundError as e:
        print(f"error: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())