giant/giant/cli.py

from collections import Counter
from datetime import datetime, timezone
from enum import Enum
import math
from pathlib import Path
import re
from typing import Optional
import uuid as uuid_mod

import numpy as np
import yaml
import torch
import typer
from typing_extensions import Annotated

import pyarrow as pa
import pyarrow.parquet as pq
from tqdm import tqdm

from giant import config as gconfig
from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
    ROLLOUT_COORD_VALUE,
)
from giant.data.loader import (
    find_parquet_files,
    iter_file_chunks,
    iter_cond_chunks,
)
from giant.data.transforms import (
    build_features,
    build_cond_features,
    decode_secondaries,
    energy_simplex_decode,
    inv_local_frame_rotation,
    inv_log_transform,
    reconstruct_post_pos,
    Normalizer,
)
from giant.geometry import GeometryOracle
from giant.model.network import build_models
from giant.particles import nearest_known_pdg
from giant.pipeline import run_train_job
from giant.rollout import rollout as run_rollout
from giant.sample import (
    sample_flow,
    sample_secondaries,
    sample_wgan,
    sample_secondaries_wgan,
)

app = typer.Typer(no_args_is_help=True)


def _router_total_experts(router_cfg: dict) -> int:
    """Total expert count for a router config, single-axis or composed.

    A composed router runs one expert per *joint* cell, so its count is the
    product of the per-axis `axis{i}_n_experts` (mirrors
    `ComposedRouter.__init__` in giant.model.network); a single-axis router
    just reports its own `n_experts`.
    """
    if router_cfg.get("type") == "composed":
        axis_counts = {
            m.group(1): int(v)
            for k, v in router_cfg.items()
            if (m := re.match(r"^axis(\d+)_n_experts$", k))
        }
        return math.prod(axis_counts.values()) if axis_counts else 1
    return int(router_cfg.get("n_experts", 1))


def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int]:
    """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.

    Routed models spend their FLOPs in the (smaller) expert trunks, not the
    monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert
    dims instead when routing is enabled. Training runs the full soft mixture
    (every expert on the whole batch), so its activation memory scales with
    the expert count; inference does top-1 dispatch (each row hits one
    expert), so the batch just partitions across experts and one expert's
    dims already bound it. estimate_batch_size scales memory linearly with
    hidden_dim * n_blocks, so the training multiplier folds into n_blocks.
    """
    router_cfg = model_cfg.get("router")
    if router_cfg and router_cfg.get("enabled"):
        hidden_dim, n_blocks = gconfig.resolve_expert_dims(
            router_cfg, model_cfg["hidden_dim"], model_cfg["n_blocks"]
        )
        if training:
            n_blocks *= _router_total_experts(router_cfg)
        return hidden_dim, n_blocks
    return model_cfg["hidden_dim"], model_cfg["n_blocks"]


def _coerce_scalar(value: str) -> object:
    """Best-effort str -> bool/int/float, else leave as str.

    CLI flag values always arrive as strings; router kwargs like
    `n_experts` (int) or `temperature` (float) need to come out typed the
    same way a TOML file's native types would, since they're merged into
    the same `model.router` dict as file-sourced config.
    """
    if value.lower() in ("true", "false"):
        return value.lower() == "true"
    try:
        return int(value)
    except ValueError:
        pass
    try:
        return float(value)
    except ValueError:
        pass
    return value


def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
    """Parse repeated `--router-axis "type:key=val,key=val"` flags into
    `axis{i}_{field}` flat keys (see `_parse_composed_axes` in
    giant.model.network), indexed by flag order — the Nth `--router-axis`
    becomes axis N.
    """
    out: dict[str, object] = {}
    for i, spec in enumerate(specs):
        axis_type, _, rest = spec.partition(":")
        out[f"axis{i}_type"] = axis_type
        for pair in filter(None, rest.split(",")):
            key, _, val = pair.partition("=")
            out[f"axis{i}_{key}"] = _coerce_scalar(val)
    return out


_CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")


def _resolve_prediction_output(data: Path, out: Path | None) -> tuple[Path, Path, str]:
    """Return (out_path, resolved_dataset_path, pred_uuid).

    When *out* is None the output path is derived from *data*:
    - under /ceph/ → fixed central store with a UUID filename
    - elsewhere   → sibling of *data* with a UUID filename
    """
    dataset_path = data.resolve()
    pred_uuid = str(uuid_mod.uuid4())
    if out is None:
        if str(dataset_path).startswith("/ceph/"):
            out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet"
        else:
            out = data.parent / f"{pred_uuid}.parquet"
    return out, dataset_path, pred_uuid


def _write_prediction_ref(
    checkpoint: Path,
    pred_uuid: str,
    out: Path,
    dataset_path: Path,
    comment: str | None = None,
) -> Path:
    """Write a YAML sidecar in the checkpoint directory and return its path."""
    ref = {
        "prediction_id": pred_uuid,
        "output": str(out),
        "dataset": str(dataset_path),
        "checkpoint": str(checkpoint.resolve()),
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    if comment is not None:
        ref["comment"] = comment
    ref_path = checkpoint.parent / f"{pred_uuid}.yaml"
    ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
    return ref_path


@app.callback()
def _main() -> None:
    """GIANT — Geant4 step-function surrogate."""


class Mode(str, Enum):
    flow = "flow"
    ddpm = "ddpm"
    wgan = "wgan"


class Conditioning(str, Enum):
    physical = "physical"
    embedding = "embedding"


class Coord(str, Enum):
    global_ = "global"
    local = "local"


class Weights(str, Enum):
    raw = "raw"
    ema = "ema"


def _load_model_weights(
    model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    ckpt: dict,
    weights: "Weights",
    checkpoint_path: Path,
) -> None:
    """Load either the raw or EMA state dicts from a training checkpoint.

    EMA weights (giant.train's shadow copy, see --ema-decay) only exist in
    checkpoints written after that feature landed, so `ema` fails loudly
    rather than silently falling back to raw weights a caller didn't ask for.
    """
    if weights == Weights.raw:
        model_key, sec_key = "model", "sec_decoder"
    else:
        model_key, sec_key = "model_ema", "sec_decoder_ema"
        if model_key not in ckpt or sec_key not in ckpt:
            typer.echo(
                f"error: {checkpoint_path} has no EMA weights (trained before "
                "--ema-decay, or with --ema-decay 0) — use --weights raw",
                err=True,
            )
            raise typer.Exit(1)
    model.load_state_dict(ckpt[model_key])
    sec_decoder.load_state_dict(ckpt[sec_key])


@app.command()
def train(
    data: Annotated[
        Path, typer.Argument(help="Parquet file or directory of parquet files")
    ],
    config: Annotated[
        Optional[Path],
        typer.Option(
            "--config", "-c", help="TOML config file (overridden by explicit flags)"
        ),
    ] = None,
    mode: Annotated[
        Optional[Mode],
        typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"),
    ] = None,
    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
    batch_size: Annotated[
        Optional[str],
        typer.Option(
            "--batch-size",
            "-b",
            help="Integer, or 'auto' to estimate from free GPU memory "
            "(cuda devices only)",
        ),
    ] = None,
    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
    weight_decay: Annotated[
        Optional[float],
        typer.Option("--weight-decay", "-W", help="AdamW weight decay (default: 0.01)"),
    ] = None,
    ema_decay: Annotated[
        Optional[float],
        typer.Option(
            "--ema-decay",
            help="EMA decay for a shadow copy of the model weights, saved "
            "alongside the raw weights in checkpoints (0 disables; default: 0.9999)",
        ),
    ] = None,
    warmup_epochs: Annotated[
        Optional[int], typer.Option("--warmup-epochs", "-w")
    ] = None,
    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
    dropout: Annotated[
        Optional[float],
        typer.Option(
            "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
        ),
    ] = None,
    conditioning: Annotated[
        Optional[Conditioning],
        typer.Option(
            "--conditioning",
            help="Input conditioning: continuous physical properties "
            "(mass/charge/Z_eff/A_eff/density/X0/lambda_int, default) or the "
            "original learned PDG/material embeddings",
        ),
    ] = None,
    router: Annotated[
        Optional[bool],
        typer.Option(
            "--router/--no-router",
            help="Route both stages through a mixture of small experts "
            "instead of one monolithic trunk (see model.router in config.toml)",
        ),
    ] = None,
    router_type: Annotated[
        Optional[str],
        typer.Option(
            "--router-type", help="Router implementation name (see ROUTER_REGISTRY)"
        ),
    ] = None,
    n_experts: Annotated[
        Optional[int], typer.Option("--n-experts", help="Number of routed experts")
    ] = None,
    router_axis: Annotated[
        Optional[list[str]],
        typer.Option(
            "--router-axis",
            help="Composed-router axis spec 'type:key=val,key=val' (repeatable; "
            "Nth flag = axis N). Use with --router-type composed instead of "
            "--n-experts, e.g. --router-axis 'energy:n_experts=4' "
            "--router-axis 'pdg:n_experts=3,emb_dim=8'",
        ),
    ] = None,
    n_critic: Annotated[
        Optional[int],
        typer.Option(
            "--n-critic",
            help="WGAN-GP (--mode wgan only): critic updates per generator "
            "update (default: 5)",
        ),
    ] = None,
    gp_weight: Annotated[
        Optional[float],
        typer.Option(
            "--gp-weight",
            help="WGAN-GP (--mode wgan only): gradient-penalty coefficient "
            "(default: 10.0)",
        ),
    ] = None,
    noise_dim: Annotated[
        Optional[int],
        typer.Option(
            "--noise-dim",
            help="WGAN (--mode wgan only): generator input noise-vector "
            "width (default: 64)",
        ),
    ] = None,
    critic_lr: Annotated[
        Optional[float],
        typer.Option(
            "--critic-lr",
            help="WGAN-GP (--mode wgan only): critic learning rate "
            "(default: same as --lr)",
        ),
    ] = None,
    val_fraction: Annotated[
        Optional[float], typer.Option("--val-fraction", "-f")
    ] = None,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", "-s", help="Random seed for reproducibility"),
    ] = None,
    validate_every: Annotated[
        Optional[int],
        typer.Option(
            "--validate-every",
            "-v",
            help="Run marginal+KL validation every N epochs (0 disables)",
        ),
    ] = None,
    validate_steps: Annotated[
        Optional[int],
        typer.Option(
            "--validate-steps",
            "-t",
            help="Flow matching ODE steps used during marginal validation "
            "(ignored in ddpm mode, which always runs the full schedule)",
        ),
    ] = None,
    max_val_batches: Annotated[
        Optional[int],
        typer.Option(
            "--max-val-batches",
            help="Cap the per-epoch val-loss pass to N batches (0 = full "
            "val set every epoch; default: 200)",
        ),
    ] = None,
    shuffle_buffer: Annotated[
        int,
        typer.Option(
            "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
        ),
    ] = 65536,
    out: Annotated[
        Optional[Path],
        typer.Option(
            "--out",
            "-o",
            help="Checkpoint dir (default: timestamped dir from hyperparams, "
            "or the --resume checkpoint's own dir when resuming)",
        ),
    ] = None,
    device: Annotated[
        Optional[str],
        typer.Option("--device", "-D", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    num_workers: Annotated[Optional[int], typer.Option("--num-workers", "-j")] = None,
    resume: Annotated[
        Optional[Path],
        typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"),
    ] = None,
    wandb: Annotated[
        Optional[bool],
        typer.Option(
            "--wandb/--no-wandb",
            help="Log per-epoch training metrics to Weights & Biases "
            "(requires `uv sync --extra wandb`)",
        ),
    ] = None,
    wandb_project: Annotated[
        Optional[str],
        typer.Option("--wandb-project", help="W&B project name (default: giant)"),
    ] = None,
    wandb_run_name: Annotated[
        Optional[str],
        typer.Option("--wandb-run-name", help="W&B run name (default: out_dir name)"),
    ] = None,
    wandb_log_every: Annotated[
        Optional[int],
        typer.Option(
            "--wandb-log-every",
            help="Log batch-level loss/grad_norm/lr to W&B every N optimizer "
            "steps (default: 50); per-epoch metrics always log in full",
        ),
    ] = None,
) -> None:
    """Train the GIANT surrogate model."""
    batch_size_auto = False
    batch_size_value: Optional[int] = None
    if batch_size is not None:
        if batch_size.strip().lower() == "auto":
            batch_size_auto = True
        else:
            try:
                batch_size_value = int(batch_size)
            except ValueError:
                typer.echo(
                    f"error: --batch-size must be an integer or 'auto', "
                    f"got {batch_size!r}",
                    err=True,
                )
                raise typer.Exit(1)

    cli_train = {
        k: v
        for k, v in {
            "mode": mode.value if mode is not None else None,
            "epochs": epochs,
            "batch_size": batch_size_value,
            "lr": lr,
            "weight_decay": weight_decay,
            "ema_decay": ema_decay,
            "warmup_epochs": warmup_epochs,
            "val_fraction": val_fraction,
            "num_workers": num_workers,
            "seed": seed,
            "validate_every": validate_every,
            "validate_steps": validate_steps,
            "max_val_batches": max_val_batches,
            "n_critic": n_critic,
            "gp_weight": gp_weight,
            "critic_lr": critic_lr,
            "wandb": wandb,
            "wandb_project": wandb_project,
            "wandb_run_name": wandb_run_name,
            "wandb_log_every": wandb_log_every,
        }.items()
        if v is not None
    }
    cli_model: dict[str, object] = {
        k: v
        for k, v in {
            "hidden_dim": hidden_dim,
            "n_blocks": n_blocks,
            "emb_dim": emb_dim,
            "dropout": dropout,
            "conditioning": conditioning.value if conditioning is not None else None,
            "noise_dim": noise_dim,
        }.items()
        if v is not None
    }
    cli_router: dict[str, object] = {
        k: v
        for k, v in {
            "enabled": router,
            "type": router_type,
            "n_experts": n_experts,
        }.items()
        if v is not None
    }
    if router_axis:
        cli_router.update(_parse_router_axis_flags(router_axis))
    if cli_router:
        cli_model["router"] = cli_router
    cfg = gconfig.merge_cli_overrides(
        gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
    )
    t, m = cfg["train"], cfg["model"]

    _device = torch.device(device) if device else gconfig.auto_device()

    if batch_size_auto:
        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m, training=True)
        try:
            t["batch_size"] = gconfig.estimate_batch_size(
                est_hidden_dim, est_n_blocks, _device
            )
        except ValueError as exc:
            typer.echo(f"error: {exc}", err=True)
            raise typer.Exit(1)
        typer.echo(
            f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
        )

    if out is not None:
        out_dir = out
    elif resume is not None:
        # Continue writing into the resumed checkpoint's own directory
        # rather than recomputing a hyperparam-derived name — the latter
        # would (a) collide with the original run's dir only by accident
        # (same day, unchanged hyperparams) and now never collides at all
        # since the fresh-run name below is timestamped to the second, and
        # (b) silently start a fresh directory if a resumed run tweaks any
        # hyperparam baked into the name (e.g. --lr for a fine-tune).
        out_dir = resume.parent
    else:
        # Name only encodes what's non-default (see default_out_dir_name), so
        # two runs with identical hyperparams in the same to-the-minute
        # timestamp would otherwise collide on this name — which also
        # doubles as the W&B run id (giant.train) — hence the suffix loop.
        base_name = gconfig.default_out_dir_name(cfg)
        out_dir = Path("checkpoints") / base_name
        suffix = 2
        while out_dir.exists():
            out_dir = Path("checkpoints") / f"{base_name}_{suffix}"
            suffix += 1

    typer.echo(f"device:  {_device}")
    typer.echo(f"out_dir: {out_dir}")

    run_train_job(
        data=data,
        cfg=cfg,
        out_dir=out_dir,
        device=_device,
        shuffle_buffer=shuffle_buffer,
        num_workers=t["num_workers"],
        resume=resume,
        echo=typer.echo,
    )


@app.command()
def predict(
    data: Annotated[
        Path, typer.Argument(help="Parquet file or directory of parquet files")
    ],
    checkpoint: Annotated[
        Path,
        typer.Option(
            "--checkpoint",
            "-c",
            help="Path to checkpoint .pt file (best.pt or last.pt)",
        ),
    ],
    coord: Annotated[
        Coord,
        typer.Option(
            "--coord",
            "-C",
            help="global: full physical units, world frame (default). "
            "local: raw 9D model output (denormalised only, local frame, "
            "log-scaled scalars) alongside the matching ground-truth target "
            "for the same input file — requires post-step columns.",
        ),
    ] = Coord.global_,
    out: Annotated[
        Optional[Path],
        typer.Option(
            "--out",
            "-o",
            help="Output parquet path (default: <data>_predicted[_local].parquet)",
        ),
    ] = None,
    batch_size: Annotated[
        str,
        typer.Option(
            "--batch-size",
            "-b",
            help="Inference batch size, or 'auto' to estimate from free GPU "
            "memory (cuda devices only)",
        ),
    ] = "4096",
    steps: Annotated[
        int,
        typer.Option(
            "--steps",
            "-s",
            help="Flow matching ODE steps (ignored for a wgan checkpoint)",
        ),
    ] = 10,
    weights: Annotated[
        Weights,
        typer.Option(
            "--weights",
            help="raw: the live training weights. ema: the EMA shadow copy "
            "(see --ema-decay in `giant train`) — usually cleaner samples, "
            "requires a checkpoint trained with EMA enabled.",
        ),
    ] = Weights.raw,
    device: Annotated[
        Optional[str],
        typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    comment: Annotated[
        Optional[str],
        typer.Option(
            "--comment",
            "-m",
            help="Free-text note recorded in the prediction's YAML sidecar",
        ),
    ] = None,
) -> None:
    """Run trained model on a parquet file and save predictions."""
    batch_size_auto = False
    batch_size_value: Optional[int] = None
    if batch_size.strip().lower() == "auto":
        batch_size_auto = True
    else:
        try:
            batch_size_value = int(batch_size)
        except ValueError:
            typer.echo(
                f"error: --batch-size must be an integer or 'auto', got {batch_size!r}",
                err=True,
            )
            raise typer.Exit(1)

    _device = torch.device(device) if device else gconfig.auto_device()
    typer.echo(f"device: {_device}")

    # --- Load checkpoint ---
    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    if "model_config" not in ckpt:
        typer.echo(
            "error: checkpoint has no model_config — retrain with the current code",
            err=True,
        )
        raise typer.Exit(1)

    if "sec_decoder" not in ckpt:
        typer.echo(
            "error: checkpoint has no sec_decoder — retrain with the current code",
            err=True,
        )
        raise typer.Exit(1)

    if "sec_phys" not in ckpt.get("normalizer", {}):
        typer.echo(
            "error: checkpoint has no normalizer.sec_phys — retrain with the "
            "current code",
            err=True,
        )
        raise typer.Exit(1)

    model_cfg = ckpt["model_config"]

    if batch_size_auto:
        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
            model_cfg, training=False
        )
        try:
            batch_size_value = gconfig.estimate_batch_size(
                est_hidden_dim,
                est_n_blocks,
                _device,
                training=False,
            )
        except ValueError as exc:
            typer.echo(f"error: {exc}", err=True)
            raise typer.Exit(1)
        typer.echo(
            f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
        )

    assert batch_size_value is not None
    bs = batch_size_value
    conditioning = model_cfg.get("conditioning", "embedding")
    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])

    model, sec_decoder = build_models(model_cfg)
    _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
    model.to(_device).eval()
    sec_decoder.to(_device).eval()

    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)

    # --- Output path ---
    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)
    typer.echo(f"output: {out}")

    # --- Stream input, generate predictions, write output ---
    files = find_parquet_files(data)
    typer.echo(f"found {len(files)} parquet file(s)")

    writer: pq.ParquetWriter | None = None
    total = 0
    skipped = 0
    unknown_pdg_counts: Counter[int] = Counter()
    total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
    chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks

    def _concat(
        a: dict[str, np.ndarray], b: dict[str, np.ndarray]
    ) -> dict[str, np.ndarray]:
        return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}

    def _process(piece: dict[str, np.ndarray]) -> None:
        nonlocal writer, total

        if coord == Coord.local:
            cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features(
                piece, pdg_map, mat_map, conditioning=conditioning
            )
            cond_cont = cond_norm.transform(cond_cont)
        else:
            cond_cont, cond_cat = build_cond_features(
                piece, pdg_map, mat_map, cond_norm, conditioning=conditioning
            )

        cc = torch.from_numpy(cond_cont).float().to(_device)
        ck = torch.from_numpy(cond_cat).long().to(_device)
        if model_cfg.get("mode") == "wgan":
            stage1_norm, n_sec_pred = sample_wgan(model, cc, ck)
        else:
            stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)

        if coord == Coord.global_:
            if model_cfg.get("mode") == "wgan":
                sec_cont, sec_phys, _sec_valid_pred = sample_secondaries_wgan(
                    sec_decoder, cc, ck, stage1_norm, n_sec_pred
                )
            else:
                sec_cont, sec_phys, _sec_valid_pred = sample_secondaries(
                    sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
                )
            sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()

        n_sec_pred_np = n_sec_pred.cpu().numpy()
        pred = stage1_norm.cpu().numpy()  # normalised

        # Inverse-normalise → local frame, log-scaled scalars
        raw = tgt_norm.inverse_transform(pred)

        if coord == Coord.local:
            table = pa.table(
                {
                    "event_id": piece["event_id"],
                    "pdg": piece["pdg"],
                    "pre_x": piece["pre_pos"][:, 0],
                    "pre_y": piece["pre_pos"][:, 1],
                    "pre_z": piece["pre_pos"][:, 2],
                    "pre_E": piece["pre_E"],
                    "pre_dx": piece["pre_dir"][:, 0],
                    "pre_dy": piece["pre_dir"][:, 1],
                    "pre_dz": piece["pre_dir"][:, 2],
                    "material": piece["material"],
                    "layer_id": piece["layer_id"],
                    "n_sec": piece["n_sec"],
                    **{
                        f"pred_{name}": raw[:, j]
                        for j, name in enumerate(LOCAL_TARGET_NAMES)
                    },
                    **{
                        f"true_{name}": target_raw[:, j]
                        for j, name in enumerate(LOCAL_TARGET_NAMES)
                    },
                }
            )
        else:
            step_length = inv_log_transform(raw[:, 0])
            # Columns 1:3 are ALR coords of the deposit/secondary/post energy
            # simplex; decode them against pre_E so edep + e_sec + post_E == pre_E
            # (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
            # doubles as the stick-breaking energy budget for the Stage-2 decode
            # below, since the model has no other source for it at inference.
            edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
                raw[:, 1:3], piece["pre_E"]
            )

            # Normalise predicted direction then rotate back to world frame
            post_dir_local = raw[:, 3:6].copy()
            norms = np.linalg.norm(post_dir_local, axis=1, keepdims=True)
            post_dir_local /= np.where(norms < 1e-8, 1.0, norms)
            post_dir_world = inv_local_frame_rotation(piece["pre_dir"], post_dir_local)

            # Same for the travel direction, then reconstruct post_pos from
            # the single shared step_length so the two stay consistent.
            travel_dir_local = raw[:, 6:9].copy()
            norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
            travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
            post_pos_world = reconstruct_post_pos(
                piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
            )

            sec_E, sec_dir_world, sec_mass, sec_charge, _sec_valid = decode_secondaries(
                sec_full_np,
                n_sec_pred_np,
                e_sec_pred,
                piece["pre_dir"],
                sec_phys_normalizer=sec_phys_norm,
            )
            # Reporting-only nearest-known-PDG label (never fed back into the
            # model) for the sec_pdg_list output column — see
            # giant/particles.py and the "no snapping at inference" design.
            sec_pdg_code = nearest_known_pdg(
                sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
            ).reshape(sec_mass.shape)
            sec_pdg_list = [
                sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
            ]
            sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
            sec_dx_list = [
                sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
            ]
            sec_dy_list = [
                sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
            ]
            sec_dz_list = [
                sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
            ]

            table = pa.table(
                {
                    "event_id": piece["event_id"],
                    "pdg": piece["pdg"],
                    "pre_x": piece["pre_pos"][:, 0],
                    "pre_y": piece["pre_pos"][:, 1],
                    "pre_z": piece["pre_pos"][:, 2],
                    "pre_E": piece["pre_E"],
                    "pre_dx": piece["pre_dir"][:, 0],
                    "pre_dy": piece["pre_dir"][:, 1],
                    "pre_dz": piece["pre_dir"][:, 2],
                    "material": piece["material"],
                    "layer_id": piece["layer_id"],
                    "n_sec": piece["n_sec"],
                    "n_sec_pred": n_sec_pred_np,
                    "step_length": step_length,
                    "delta_e": delta_e,
                    "edep": edep,
                    "post_dx": post_dir_world[:, 0],
                    "post_dy": post_dir_world[:, 1],
                    "post_dz": post_dir_world[:, 2],
                    "post_x": post_pos_world[:, 0],
                    "post_y": post_pos_world[:, 1],
                    "post_z": post_pos_world[:, 2],
                    "sec_pdg_list": sec_pdg_list,
                    "sec_E_list": sec_E_list,
                    "sec_dx_list": sec_dx_list,
                    "sec_dy_list": sec_dy_list,
                    "sec_dz_list": sec_dz_list,
                }
            )

        table = table.replace_schema_metadata(
            {
                PREDICT_COORD_METADATA_KEY: coord.value,
                PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
            }
        )

        if writer is None:
            writer = pq.ParquetWriter(out, table.schema)
        writer.write_table(table)
        total += len(piece["event_id"])

    # Buffer rows across row-group boundaries so the inference batch size
    # isn't capped by however the source file happens to be chunked.
    buffer: dict[str, np.ndarray] | None = None

    bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
    for path in files:
        for chunk in chunk_iter(path):
            N_in = len(chunk["event_id"])

            pdg_mask = np.array([int(p) in pdg_map for p in chunk["pdg"]])
            if not pdg_mask.all():
                unknown_pdg_counts.update(int(p) for p in chunk["pdg"][~pdg_mask])
                chunk = {k: v[pdg_mask] for k, v in chunk.items()}

            skipped += N_in - len(chunk["event_id"])
            bar.update(N_in)
            if len(chunk["event_id"]) == 0:
                continue

            buffer = chunk if buffer is None else _concat(buffer, chunk)
            while len(buffer["event_id"]) >= bs:
                piece = {k: v[:bs] for k, v in buffer.items()}
                buffer = {k: v[bs:] for k, v in buffer.items()}
                _process(piece)

    if buffer is not None and len(buffer["event_id"]) > 0:
        _process(buffer)

    bar.close()
    if writer is not None:
        writer.close()

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment)
    typer.echo(f"reference: {ref_path}")

    if skipped:
        codes = ", ".join(
            f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())
        )
        typer.echo(
            f"warning: skipped {skipped:,} row(s) with unknown PDG code(s): {codes}",
            err=True,
        )
    typer.echo(f"wrote {total:,} rows → {out}")


def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.ndarray]:
    """Pick each event's primary entry state (argmax-pre_E row) as a shower seed.

    Streams conditioning columns and keeps the highest-pre_E step per event_id —
    the codebase's convention for the primary (a secondary always carries less
    energy than its parent). See giant/analysis/reduce.py:entry_axis.
    """
    best_E: dict[int, float] = {}
    best: dict[int, tuple] = {}
    for path in files:
        for chunk in iter_cond_chunks(path):
            ev = chunk["event_id"]
            pe = chunk["pre_E"]
            for i in range(len(ev)):
                e = int(ev[i])
                if pe[i] > best_E.get(e, -np.inf):
                    best_E[e] = float(pe[i])
                    best[e] = (
                        int(chunk["pdg"][i]),
                        chunk["pre_pos"][i].astype(np.float64),
                        float(pe[i]),
                        chunk["pre_dir"][i].astype(np.float64),
                    )
    event_ids = sorted(best)
    if n_events is not None:
        event_ids = event_ids[:n_events]
    if not event_ids:
        raise ValueError("no events found to seed from")

    return {
        "event_id": np.array(event_ids, dtype=np.int64),
        "pdg": np.array([best[e][0] for e in event_ids], dtype=np.int64),
        "pre_pos": np.stack([best[e][1] for e in event_ids]),
        "pre_E": np.array([best[e][2] for e in event_ids], dtype=np.float64),
        "pre_dir": np.stack([best[e][3] for e in event_ids]),
    }


@app.command()
def rollout(
    data: Annotated[
        Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")
    ],
    checkpoint: Annotated[
        Path,
        typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"),
    ],
    geometry: Annotated[
        Path,
        typer.Option(
            "--geometry",
            "-g",
            help="Geometry oracle .pkl (dwarf build-geometry-oracle)",
        ),
    ],
    energy_cutoff: Annotated[
        float,
        typer.Option(
            "--energy-cutoff",
            help="Stop a track when its energy drops below this [MeV]",
        ),
    ] = 0.1,
    max_steps: Annotated[
        int, typer.Option("--max-steps", help="Max steps per individual track")
    ] = 1000,
    steps: Annotated[
        int,
        typer.Option(
            "--steps",
            "-s",
            help="Flow matching ODE steps per model call (ignored for a wgan checkpoint)",
        ),
    ] = 10,
    weights: Annotated[
        Weights,
        typer.Option(
            "--weights",
            help="raw: the live training weights. ema: the EMA shadow copy "
            "(see --ema-decay in `giant train`) — usually cleaner samples, "
            "requires a checkpoint trained with EMA enabled.",
        ),
    ] = Weights.raw,
    batch_size: Annotated[
        int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
    ] = 4096,
    max_tracks_per_event: Annotated[
        Optional[int],
        typer.Option(
            "--max-tracks-per-event",
            help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
        ),
    ] = None,
    escape_threshold: Annotated[
        Optional[float],
        typer.Option(
            "--escape-threshold",
            help="Override the oracle's NN-distance escape threshold [mm]",
        ),
    ] = None,
    n_events: Annotated[
        Optional[int], typer.Option("--n-events", help="Cap number of seed events")
    ] = None,
    device: Annotated[
        Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")
    ] = None,
    out: Annotated[
        Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
    ] = None,
    seed: Annotated[
        Optional[int],
        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
    ] = None,
) -> None:
    """Roll the surrogate forward into full showers (autoregressive)."""
    if seed is not None:
        torch.manual_seed(seed)
        np.random.seed(seed)

    _device = torch.device(device) if device else gconfig.auto_device()
    typer.echo(f"device: {_device}")

    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    for key in ("model_config", "sec_decoder"):
        if key not in ckpt:
            typer.echo(
                f"error: checkpoint has no {key} — retrain with the current code",
                err=True,
            )
            raise typer.Exit(1)

    if "sec_phys" not in ckpt.get("normalizer", {}):
        typer.echo(
            "error: checkpoint has no normalizer.sec_phys — retrain with the "
            "current code",
            err=True,
        )
        raise typer.Exit(1)

    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
    training_cfg = gconfig.load_checkpoint_config(checkpoint)

    model_cfg = ckpt["model_config"]
    conditioning = model_cfg.get("conditioning", "embedding")
    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])

    model, sec_decoder = build_models(model_cfg)
    _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
    model.to(_device).eval()
    sec_decoder.to(_device).eval()
    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")

    oracle = GeometryOracle.load(geometry)
    typer.echo(
        f"loaded geometry oracle: {geometry} "
        f"(escape_threshold={oracle.escape_threshold:.3f})"
    )

    files = find_parquet_files(data)
    seeds = _seed_from_data(files, n_events)
    typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")

    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)

    # Written incrementally as each batch of steps is produced, rather than
    # buffering the whole run (which scales with n_events * max_steps *
    # avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
    # already does on its input side.
    writer: pq.ParquetWriter | None = None

    def _write_chunk(row: dict[str, np.ndarray]) -> None:
        nonlocal writer
        table = pa.table(row)
        if writer is None:
            table = table.replace_schema_metadata(
                {
                    PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE,
                    PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
                }
            )
            writer = pq.ParquetWriter(out, table.schema)
        writer.write_table(table)

    summary = run_rollout(
        model,
        sec_decoder,
        oracle,
        seeds,
        cond_norm,
        tgt_norm,
        sec_phys_norm,
        pdg_map,
        mat_map,
        energy_cutoff=energy_cutoff,
        max_steps=max_steps,
        steps=steps,
        batch_size=batch_size,
        device=_device,
        max_tracks_per_event=max_tracks_per_event,
        escape_threshold=escape_threshold,
        on_chunk=_write_chunk,
        conditioning=conditioning,
        mode=model_cfg.get("mode", "flow"),
    )
    if writer is not None:
        writer.close()

    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
    ref = yaml.safe_load(ref_path.read_text())
    ref.update(
        {
            "kind": "rollout",
            "geometry_oracle": str(geometry.resolve()),
            "energy_cutoff": energy_cutoff,
            "max_steps": max_steps,
            "steps": steps,
            "max_tracks_per_event": max_tracks_per_event,
            "escape_threshold": escape_threshold,
            "n_events": n_events,
            "n_seed_events": int(len(seeds["event_id"])),
            "weights": weights.value,
            "batch_size": batch_size,
            "device": str(_device),
            "rollout_seed": seed,
            "n_rows": summary["n_rows"],
            "termination_reason_counts": summary["termination_reason_counts"],
            # Full architecture spec baked into the checkpoint — includes the
            # entire router sub-dict, not just a hand-picked subset, so any
            # model knob (router type/n_experts, noise_dim, vocab sizes, ...)
            # is available downstream without touching this command again.
            "model_config": dict(model_cfg),
            "training_epoch": ckpt.get("epoch"),
            "best_val_loss": ckpt.get("best_val_loss"),
            # [train]/[meta] from the sibling config.toml (giant.config.save_config)
            # — empty dicts if the checkpoint has no config.toml next to it.
            "training_config": dict(training_cfg.get("train", {})),
            "training_meta": dict(training_cfg.get("meta", {})),
        }
    )
    ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))

    typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
    typer.echo(f"terminations: {summary['termination_reason_counts']}")
    typer.echo(f"reference: {ref_path}")


analyze_app = typer.Typer(
    no_args_is_help=True,
    help="Rollout-vs-reference analysis: parallel compute on HTCondor + local render.",
)
app.add_typer(analyze_app, name="analyze")


@analyze_app.command("prep")
def analyze_prep(
    rollout_yaml: Annotated[
        Path,
        typer.Argument(
            help="giant rollout YAML sidecar (names the rollout + reference files)"
        ),
    ],
    run_dir: Annotated[
        Optional[Path],
        typer.Option(
            "--run-dir",
            "-o",
            help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
        ),
    ] = None,
    n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
    n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
    top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
    chunks: Annotated[
        int,
        typer.Option(
            "--chunks", help="Split each plot's data into this many event_id chunks"
        ),
    ] = 1,
) -> None:
    """Read the rollout YAML → shared.json + run_meta.json in the run directory."""
    from giant.analysis import prep

    path = prep(
        rollout_yaml,
        run_dir,
        n_chunks=chunks,
        default_base=Path.cwd() / "analysis_runs",
        n_energy_bins=n_energy_bins,
        n_marginal_bins=n_marginal_bins,
        top_k_pdg=top_k_pdg,
    )
    typer.echo(f"run directory: {path}")


@analyze_app.command("compute-one")
def analyze_compute_one(
    id: Annotated[
        str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
    ],
    run_dir: Annotated[
        Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
    ],
    chunk: Annotated[
        int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")
    ] = 0,
) -> None:
    """Run one (plot, chunk)'s streaming reduction (this is what each condor job runs)."""
    from giant.analysis import compute_one

    path = compute_one(id, run_dir, chunk_index=chunk)
    typer.echo(f"wrote {path}")


@analyze_app.command("merge-one")
def analyze_merge_one(
    id: Annotated[
        str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
    ],
    run_dir: Annotated[
        Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
    ],
) -> None:
    """Merge one plot's chunk partials into its final reduced JSON.

    Runs automatically as part of `analyze render`; useful standalone to
    debug a specific plot without re-rendering everything.
    """
    from giant.analysis import merge_one

    path = merge_one(id, run_dir)
    typer.echo(f"wrote {path}")


@analyze_app.command("list")
def analyze_list() -> None:
    """Print every catalog plot id."""
    from giant.analysis import catalog_ids

    for pid in catalog_ids():
        typer.echo(pid)


@analyze_app.command("render")
def analyze_render(
    run_dir: Annotated[
        Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")
    ],
    gallery: Annotated[
        bool,
        typer.Option(
            "--gallery/--no-gallery", help="Run `gallery generate` after rendering"
        ),
    ] = False,
) -> None:
    """Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX)."""
    from giant.analysis.render import render_run

    pdfs = render_run(run_dir, run_gallery=gallery)
    typer.echo(f"rendered {len(pdfs)} plots → {Path(run_dir) / 'plots'}")


@analyze_app.command("submit")
def analyze_submit(
    rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
    accounting_group: Annotated[str, typer.Option("--accounting-group")],
    run_dir: Annotated[
        Optional[Path],
        typer.Option(
            "--run-dir",
            "-o",
            help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
        ),
    ] = None,
    docker_image: Annotated[
        str, typer.Option("--docker-image")
    ] = "cverstege/alma9-gridjob",
    request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192,
    remote: Annotated[
        bool,
        typer.Option("--remote/--local", help="+RemoteJob vs ProvidesETPResources"),
    ] = False,
    chunks: Annotated[
        int,
        typer.Option(
            "--chunks",
            help="Split each plot's data into this many event_id chunks/jobs",
        ),
    ] = 1,
    n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
    n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
    top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
    dry_run: Annotated[
        bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
    ] = False,
) -> None:
    """prep + write the HTCondor submit description (one job per plot x chunk), then submit."""
    import subprocess

    from giant.analysis import SubmitConfig, prep, write_submit

    path = prep(
        rollout_yaml,
        run_dir,
        n_chunks=chunks,
        default_base=Path.cwd() / "analysis_runs",
        n_energy_bins=n_energy_bins,
        n_marginal_bins=n_marginal_bins,
        top_k_pdg=top_k_pdg,
    )
    cfg = SubmitConfig(
        run_dir=path,
        accounting_group=accounting_group,
        repo_dir=Path.cwd(),
        docker_image=docker_image,
        request_memory_mb=request_memory,
        remote=remote,
        n_chunks=chunks,
    )
    sub = write_submit(cfg)
    typer.echo(f"run directory: {path}")
    typer.echo(f"wrote submit description: {sub}")
    if dry_run:
        typer.echo("dry-run: not submitting")
        return
    subprocess.run(["condor_submit", str(sub)], check=True)


if __name__ == "__main__":
    app()