onf-v3/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) and ``train`` (the retrieval-network curriculum, see
:mod:`onf.graph.train`).

This module used to also carry ``probe`` (the falsification battery), ``eval-offline`` (rung R0 of the
evaluation ladder), ``probe-head`` (a trained checkpoint's abstain-decision diagnostic), ``report`` and
``list`` (run-bookkeeping inspection). All five are now CLOSED ABLATIONS / retired tooling and have been
deleted along with the helpers that only served them (``onf.graph.probes`` -- their sole non-CLI
dependency -- was deleted first; the three subcommands importing it already raised
``ModuleNotFoundError`` before this cut). The repo's supported CLI surface is exactly ``build`` and
``train``.

WHY A CLI, NOT A PILE OF SCRIPTS
    Both library functions this module calls already do the real work (:func:`onf.graph.builders.
    build_graph`, :func:`onf.graph.train.train_graph`) and already accept a :class:`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``
    :meth:`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 :class:`StageLogger` repoints it (the default). ``train`` writes its checkpoint straight
    into that SAME resolved directory (mirroring :func:`onf.graph.train.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 :class:`StageLogger` -- a
    `train` run would otherwise silently orphan the graph it just extended from ``Paths.graph``'s point
    of view.

EXIT CODES (every subcommand)
    0 success -- 1 an exception / bad input. Neither subcommand has a GATE (exit code 2 is unused by
    this module now; it belonged to the deleted probe/eval-offline/probe-head gates).
"""

from __future__ import annotations

import argparse
import dataclasses
import glob
import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
import yaml

from onf.config import GraphConfig, default_paths
from onf.graph import schema
from onf.graph.builders import _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.builders.build_graph /
#: onf.graph.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)


class CLIError(Exception):
    """Raised by any subcommand to fail cleanly with a specific exit code (1 by default) instead of a
    raw traceback -- :func:`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: the ``suites:`` mapping keys of
    ``configs/suites.yaml``, 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."""
    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 :meth:`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 _csv(value: str) -> list[str]:
    return [v.strip() for v in value.split(",") if v.strip()]


def _csv_int(value: str) -> list[int]:
    return [int(v) for v in _csv(value)]


def _json_default(obj: Any) -> Any:
    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(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:
    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)


# ======================================================================================================
# outputs/ run-directory scanning (the generic "skip if already done" check every subcommand uses)
# ======================================================================================================
@dataclasses.dataclass
class _RunInfo:
    suite: str
    stage: str
    run_dir: Path
    manifest: dict
    metrics: dict
    config: dict


def _load_run(run_dir: Path) -> _RunInfo | None:
    manifest_p, metrics_p, config_p = (run_dir / n for n in ("manifest.json", "metrics.json", "config.json"))
    if not manifest_p.exists():
        return None
    try:
        manifest = json.loads(manifest_p.read_text())
    except (json.JSONDecodeError, OSError):
        return None
    metrics = {}
    if metrics_p.exists():
        try:
            metrics = json.loads(metrics_p.read_text())
        except (json.JSONDecodeError, OSError):
            metrics = {}
    config = {}
    if config_p.exists():
        try:
            config = json.loads(config_p.read_text())
        except (json.JSONDecodeError, OSError):
            config = {}
    return _RunInfo(
        suite=str(manifest.get("suite", run_dir.parent.name)),
        stage=str(manifest.get("stage", run_dir.name.rsplit("_", 1)[0])),
        run_dir=run_dir, manifest=manifest, metrics=metrics, config=config,
    )


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: the
    newest COMPLETE run under ``outputs/<suite>/<stage>_*`` whose ``config.json`` matches ``config``
    (ignoring the purely cosmetic ``name`` key -- two runs with the same knobs and different
    ``--name`` labels are still the same work)."""
    base = root / suite
    if not base.is_dir():
        return None
    norm_target = {k: v for k, v in config.items() if k != "name"}
    candidates = sorted(base.glob(f"{stage}_*"), reverse=True)
    for run_dir in candidates:
        if not run_dir.is_dir() or run_dir.is_symlink():
            continue
        info = _load_run(run_dir)
        if info is None or info.manifest.get("status") != "complete":
            continue
        norm_got = {k: v for k, v in info.config.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:
    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?)")
    try:
        return json.loads(p.read_text())
    except (json.JSONDecodeError, OSError) as e:
        raise CLIError(f"--resume {run_dir}: unreadable config.json ({e})") from e


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
    ``--resume``d 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


# ======================================================================================================
# common argparse plumbing shared by every subcommand
# ======================================================================================================
def _add_common_args(sub: argparse.ArgumentParser, *, resume: bool = True) -> None:
    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/")
    if resume:
        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 _stage_logger(args: argparse.Namespace, stage: str, suite: str, config: dict, *,
                   n_stages: int = 0, inputs: list[str] | None = None, update_latest: bool = True) -> StageLogger:
    return StageLogger(
        stage, suite, name=args.name, root=args.out, config=config, n_stages=n_stages,
        inputs=inputs or [], quiet=args.quiet, update_latest=update_latest,
    )


# ======================================================================================================
# build
# ======================================================================================================
def _add_build_parser(sub: argparse._SubParsersAction) -> None:
    p = 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.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.",
    )
    p.add_argument("--suite", required=True, help="object | spatial | goal | long | all "
                    "(configs/suites.yaml)")
    p.add_argument("--coarsen", type=int, default=None, help=f"raw frames per node "
                    f"(default: schema.COARSEN={schema.COARSEN})")
    p.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)")
    p.add_argument("--limit-demos", type=int, default=None, help="cap demos read per hdf5 file "
                    "(smoke-test a suite without the full 500 demos)")
    p.add_argument("--limit-tasks", type=int, default=None, help="cap hdf5 files (tasks) read")
    _add_common_args(p)
    p.set_defaults(_run=_cmd_build)


def _build_one_suite(args: argparse.Namespace, suite: str, paths, resumed: dict) -> dict:
    cfg = GraphConfig(
        coarsen=_pick(args.coarsen, resumed, "coarsen", schema.COARSEN),
        device=_pick(args.device, resumed, "device", ""),
    )
    hdf5_dir = _pick(args.hdf5_dir, resumed, "hdf5_dir", None)
    hdf5_dir = os.fspath(hdf5_dir) if hdf5_dir else os.fspath(paths.hdf5(_hdf5_name(suite)))
    seed = _pick(args.seed, resumed, "seed", 0)
    limits = {
        k: v for k, v in (
            ("limit_demos", _pick(args.limit_demos, resumed, "limit_demos", None)),
            ("limit_tasks", _pick(args.limit_tasks, resumed, "limit_tasks", None)),
        ) if v is not None
    }
    config = {
        "command": "build", "suite": suite, "hdf5_dir": hdf5_dir, "seed": seed, "name": args.name,
        "limits": limits, **dataclasses.asdict(cfg),
    }

    if args.dry_run:
        _echo_config(f"build[{suite}]", config)
        print(f"planned stages ({len(BUILD_STAGE_LABELS)}): {', '.join(BUILD_STAGE_LABELS)}")
        return {"suite": suite, "dry_run": True}

    root = Path(args.out) if args.out else paths.outputs()
    if not args.force:
        hit = _find_matching_run(root, suite, "graph", config)
        if hit is not None:
            print(f"[build] {suite}: SKIP -- identical config already built at {hit} "
                  f"(pass --force to rebuild)")
            return {"suite": suite, "skipped": True, "run_dir": str(hit)}

    _echo_config(f"build[{suite}]", config)
    input_files = sorted(glob.glob(f"{hdf5_dir}/*.hdf5"))
    with _stage_logger(args, "graph", suite, config, n_stages=len(BUILD_STAGE_LABELS),
                        inputs=input_files, update_latest=True) as log:
        metrics = build_graph(
            suite, hdf5_dir=hdf5_dir, out_dir=None, cfg=cfg, coarsen=cfg.coarsen,
            device=(cfg.device or None), logger=log, **limits,
        )
    return {"suite": suite, "skipped": False, **metrics}


def _cmd_build(args: argparse.Namespace) -> int:
    paths = default_paths()
    configured = load_configured_suites()
    suites = _resolve_suites(args.suite, configured)
    resumed = _load_resume_config(args.resume) if getattr(args, "resume", None) else {}

    ok, results = True, []
    for suite in suites:
        try:
            results.append(_build_one_suite(args, suite, paths, resumed))
        except Exception as e:
            ok = False
            print(f"[build] {suite}: FAILED -- {type(e).__name__}: {e}", file=sys.stderr)

    if args.dry_run:
        return 0
    rows = [
        [r["suite"], "skip" if r.get("skipped") else "built",
         str(r.get("n_nodes", "-")), str(r.get("n_edges", "-")), str(r.get("n_stages", "-"))]
        for r in results
    ]
    _print_table(["suite", "status", "n_nodes", "n_edges", "n_stages"], rows)
    return 0 if ok else 1


# ======================================================================================================
# train
# ======================================================================================================
def _add_train_parser(sub: argparse._SubParsersAction) -> None:
    p = sub.add_parser(
        "train", help="train GraphRetrieverNet (onf.graph.train) on a built graph",
        description="Run the demo-graph retrieval curriculum (stage 1, temporal-advance pretraining -- "
        "the only stage left; stage 2's drift finetuning was a closed ablation and has been deleted, "
        "see onf.graph.train's module docstring) and write g_head.npz next to the graph it was trained "
        "on. Requires `build` to have run first.",
    )
    p.add_argument("--suite", required=True, help="object | spatial | goal | long | all")
    p.add_argument("--epochs", type=int, default=None, help=f"epochs per stage "
                    f"(default: {DEFAULT_TRAIN_EPOCHS})")
    p.add_argument("--n-query", type=int, default=None, help="training queries per stage "
                    "(default: onf.graph.train.N_QUERY)")
    p.add_argument("--n-eval", type=int, default=None, help="held-out queries for per-epoch metrics "
                    "(default: onf.graph.train.N_EVAL)")
    p.add_argument("--graph-dir", default=None, help="override the graph dir to train on/into "
                    "(default: Paths.graph(<suite>))")
    _add_common_args(p)
    p.set_defaults(_run=_cmd_train)


def _train_one_suite(args: argparse.Namespace, suite: str, paths, resumed: dict) -> dict:
    from onf.graph.train import N_EVAL, N_QUERY

    stages = (1,)   # only legal value -- see onf.graph.train.train_graph
    epochs = _pick(args.epochs, resumed, "epochs", DEFAULT_TRAIN_EPOCHS)
    n_query = _pick(args.n_query, resumed, "n_query", N_QUERY)
    n_eval = _pick(args.n_eval, resumed, "n_eval", N_EVAL)
    seed = _pick(args.seed, resumed, "seed", 0)
    device = _pick(args.device, resumed, "device", "")
    graph_dir_arg = _pick(args.graph_dir, resumed, "graph_dir", None)
    graph_dir = Path(graph_dir_arg) if graph_dir_arg else Path(paths.graph(suite))

    cfg = GraphConfig(device=device)
    config = {
        "command": "train", "suite": suite, "graph_dir": str(graph_dir), "stages": list(stages),
        "epochs": epochs, "n_query": n_query, "n_eval": n_eval, "seed": seed, "name": args.name,
        **dataclasses.asdict(cfg),
    }

    if args.dry_run:
        _echo_config(f"train[{suite}]", config)
        print(f"planned stages: {stages} (graph_dir={graph_dir})")
        return {"suite": suite, "dry_run": True}

    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")

    root = Path(args.out) if args.out else paths.outputs()
    checkpoint_path = graph_dir / schema.HEAD_NPZ
    if not args.force:
        hit = _find_matching_run(root, suite, "train", config)
        if hit is not None and checkpoint_path.exists():
            print(f"[train] {suite}: SKIP -- identical config already trained at {hit}, "
                  f"checkpoint={checkpoint_path} (pass --force to retrain)")
            return {"suite": suite, "skipped": True, "run_dir": str(hit)}

    _echo_config(f"train[{suite}]", config)
    from onf.graph.train import train_graph

    with _stage_logger(args, "train", suite, config, n_stages=len(stages),
                        inputs=[str(graph_dir / schema.EDGES_NPZ)], update_latest=False) as log:
        result = train_graph(
            graph_dir, stages=stages, epochs=epochs, seed=seed, device=(device or None), save=True,
            cfg=cfg, logger=log, n_query=n_query, n_eval=n_eval, suite=suite,
        )
        log.metric(checkpoint_path=result["checkpoint_path"], graph_hash=result["graph_hash"])
    final_metrics = {}
    for stage_hist in result["history"].values():
        if stage_hist and isinstance(stage_hist[-1], dict) and stage_hist[-1].get("epoch") == "final":
            final_metrics = stage_hist[-1]
    return {"suite": suite, "skipped": False, "checkpoint_path": result["checkpoint_path"],
            "final": final_metrics}


def _cmd_train(args: argparse.Namespace) -> int:
    paths = default_paths()
    configured = load_configured_suites()
    suites = _resolve_suites(args.suite, configured)
    resumed = _load_resume_config(args.resume) if getattr(args, "resume", None) else {}

    ok, results = True, []
    for suite in suites:
        try:
            results.append(_train_one_suite(args, suite, paths, resumed))
        except CLIError:
            raise
        except Exception as e:
            ok = False
            print(f"[train] {suite}: FAILED -- {type(e).__name__}: {e}", file=sys.stderr)

    if args.dry_run:
        return 0
    rows = []
    for r in results:
        final = r.get("final", {})
        rows.append([
            r["suite"], "skip" if r.get("skipped") else "trained",
            _fmt_metric(final.get("top1", "-")), _fmt_metric(final.get("mrr", "-")),
            _fmt_metric(final.get("crossed_top1", "-")),
        ])
    _print_table(["suite", "status", "top1", "mrr", "crossed_top1"], rows)
    return 0 if ok else 1


# ======================================================================================================
# top-level parser / dispatch
# ======================================================================================================
def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="python -m onf.graph",
        description="The demonstration-graph subsystem's offline pipeline (see onf.graph.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")
    _add_build_parser(sub)
    _add_train_parser(sub)
    return p


def main(argv: list[str] | None = None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return int(args._run(args))
    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())