from __future__ import annotations

import math
import re
import uuid as uuid_mod
from collections import Counter
from datetime import UTC, datetime
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, cast

import typer

if TYPE_CHECKING:
    import numpy as np

from giant import config as gconfig
from giant.constants import (
    LOCAL_TARGET_NAMES,
    PREDICT_COORD_METADATA_KEY,
    PREDICT_SCHEMA_VERSION,
    PREDICT_SCHEMA_VERSION_KEY,
    PREDICT_TRUTH_METADATA_KEY,
    ROLLOUT_COORD_VALUE,
)

# giant.materials only pulls in numpy (no torch/pandas), and MATERIAL_PROPERTIES
# is needed at decoration time below (a Typer option default), so it can't be
# deferred into a command body like the rest of this module's heavy imports.
from giant.materials import MATERIAL_PROPERTIES

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, stage: str = "stage1") -> tuple[int, int]:
    """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.

    `model_cfg` is either the new nested shape (has a `f"{stage}_model"` key
    — the merged training `cfg`, or a checkpoint's new-format `model_config`)
    or a v0.2 checkpoint's flat `model_config`. v0.3.0 dropped per-expert
    sizing (giant.model.network's routed trunks always inherit the stage's
    own hidden_dim/n_res_blocks — no more `resolve_expert_dims`), so the new
    shape needs no special-casing there; the legacy flat shape may still
    carry a v0.2 `expert_hidden_dim`/`expert_n_blocks` override, honoured
    only when that checkpoint's router was actually enabled.

    Routed models spend their FLOPs in the (smaller) expert trunks. 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.
    """
    if f"{stage}_model" in model_cfg:
        stage_cfg = model_cfg[f"{stage}_model"]
        hidden_dim, n_blocks = stage_cfg["hidden_dim"], stage_cfg["n_res_blocks"]
        router_cfg = stage_cfg.get("router")
    else:
        hidden_dim, n_blocks = model_cfg["hidden_dim"], model_cfg["n_blocks"]
        router_cfg = model_cfg.get("router")
        if router_cfg and router_cfg.get("enabled"):
            hidden_dim = model_cfg.get("expert_hidden_dim") or hidden_dim
            n_blocks = model_cfg.get("expert_n_blocks") or n_blocks
    if router_cfg and router_cfg.get("enabled") and training:
        n_blocks = n_blocks * _router_total_experts(router_cfg)
    return hidden_dim, 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


def _parse_set_flags(specs: list[str] | None) -> dict[str, object]:
    """Parse repeated `--set dotted.path=value` flags into a dict, typing
    each value with `_coerce_scalar` the same way a TOML file's native types
    would arrive. Validation against the inference-safe allowlist happens
    downstream in `giant.checkpoint_io.apply_config_overrides` — this only
    parses syntax.
    """
    out: dict[str, object] = {}
    for spec in specs or []:
        path, sep, val = spec.partition("=")
        if not sep:
            typer.echo(f"error: --set {spec!r} must be 'dotted.path=value'", err=True)
            raise typer.Exit(1)
        out[path] = _coerce_scalar(val)
    return out


def _router_cli_overrides(
    router: bool | None,
    router_type: str | None,
    n_experts: int | None,
    router_axis: list[str] | None,
) -> dict[str, object]:
    """Build the `model.router` override dict from `--router`/`--router-type`/
    `--n-experts`/`--router-axis` flags (empty if none were given). Shared by
    `train` and `new-run` so both resolve router overrides identically.
    """
    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))
    return cli_router


_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,
    extra: dict | None = None,
) -> Path:
    """Write a YAML sidecar in the checkpoint directory and return its path.

    ``extra`` is merged in after the base fields (e.g. `giant rollout`'s
    provenance/timing block, or `giant predict`'s) — callers that don't pass
    it get exactly today's thin sidecar.
    """
    import yaml

    ref = {
        "kind": "prediction",
        "prediction_id": pred_uuid,
        "output": str(out),
        "dataset": str(dataset_path),
        "checkpoint": str(checkpoint.resolve()),
        "timestamp": datetime.now(UTC).isoformat(),
    }
    if comment is not None:
        ref["comment"] = comment
    if extra:
        ref.update(extra)
    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


def _build_rollout_timing(
    *,
    setup_s: float,
    rollout_s: float,
    write_s: float,
    n_rows: int,
    termination_reason_counts: dict[str, int],
    n_seed_events: int,
    device: str,
    torch_threads: int,
) -> dict:
    """Assemble ``giant rollout``'s ``timing`` sidecar block.

    ``n_physical_rows`` excludes the synthetic termination rows (escape/
    unknown-pdg/energy-cutoff/max-steps markers `giant.rollout` emits but
    Geant4 never does) so ``us_per_step`` is comparable to
    ``giant.analysis.geant4_reference``'s per-step Geant4 measurement — see
    ``giant/analysis/catalog.py``'s ``eval_cost_per_step`` spec.
    """
    from giant.analysis.sources import SYNTHETIC_TERMINATION_REASONS

    sample_s = rollout_s - write_s
    n_synthetic_rows = sum(termination_reason_counts.get(reason, 0) for reason in SYNTHETIC_TERMINATION_REASONS)
    n_physical_rows = n_rows - n_synthetic_rows
    return {
        "setup_s": setup_s,
        "rollout_s": rollout_s,
        "write_s": write_s,
        "sample_s": sample_s,
        "n_rows": n_rows,
        "n_physical_rows": n_physical_rows,
        "us_per_step": (sample_s / n_physical_rows * 1e6) if n_physical_rows else None,
        "write_us_per_step": (write_s / n_physical_rows * 1e6) if n_physical_rows else None,
        "ms_per_event": (rollout_s / n_seed_events * 1e3) if n_seed_events else None,
        "device": device,
        "torch_threads": torch_threads,
    }


def _build_predict_timing(
    *,
    setup_s: float,
    predict_s: float,
    write_s: float,
    n_rows: int,
    device: str,
    torch_threads: int,
) -> dict:
    """Assemble ``giant predict``'s ``timing`` sidecar block.

    Keys are deliberately compatible with ``_build_rollout_timing``'s (same
    names for the quantities both commands have) so a gallery's ``timing``
    metadata renders the same way whether the series came from a rollout or
    a prediction. There's no ``n_physical_rows``/``ms_per_event`` here:
    unlike a rollout, `giant predict` never emits synthetic termination rows
    (one output row per input step) and doesn't work in whole showers/events
    — so ``us_per_step`` is already directly comparable to a rollout's and to
    ``giant.analysis.geant4_reference``'s per-step Geant4 measurement.
    """
    sample_s = predict_s - write_s
    return {
        "setup_s": setup_s,
        "predict_s": predict_s,
        "write_s": write_s,
        "sample_s": sample_s,
        "n_rows": n_rows,
        "us_per_step": (sample_s / n_rows * 1e6) if n_rows else None,
        "write_us_per_step": (write_s / n_rows * 1e6) if n_rows else None,
        "rows_per_s": (n_rows / predict_s) if predict_s else None,
        "device": device,
        "torch_threads": torch_threads,
    }


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


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


class Decoder(str, Enum):
    one_shot = "one_shot"
    autoregressive = "autoregressive"


class Stage1Context(str, Enum):
    truth = "truth"
    sampled = "sampled"


# Conditioning itself lives in giant.config (imported below as gconfig) —
# shared with giant/tools/dwarf.py's Typer commands so the two CLIs can't
# silently drift apart on the option's valid values.
Conditioning = gconfig.Conditioning


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


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


@app.command()
def train(
    data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
    config: Annotated[
        Path | None,
        typer.Option("--config", "-c", help="TOML config file (overridden by explicit flags)"),
    ] = None,
    mode: Annotated[
        Mode | None,
        typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"),
    ] = None,
    epochs: Annotated[int | None, typer.Option("--epochs", "-e")] = None,
    batch_size: Annotated[
        str | None,
        typer.Option(
            "--batch-size",
            "-b",
            help="Integer, or 'auto' to estimate from free GPU memory (cuda devices only)",
        ),
    ] = None,
    lr: Annotated[float | None, typer.Option("--lr", "-l")] = None,
    weight_decay: Annotated[
        float | None,
        typer.Option("--weight-decay", "-W", help="AdamW weight decay (default: 0.01)"),
    ] = None,
    ema_decay: Annotated[
        float | None,
        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[int | None, typer.Option("--warmup-epochs", "-w")] = None,
    hidden_dim: Annotated[int | None, typer.Option("--hidden-dim", "-H")] = None,
    n_blocks: Annotated[int | None, typer.Option("--n-blocks", "-n")] = None,
    emb_dim: Annotated[int | None, typer.Option("--emb-dim", "-E")] = None,
    dropout: Annotated[
        float | None,
        typer.Option("--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"),
    ] = None,
    stage1_generator: Annotated[
        Mode | None,
        typer.Option(
            "--stage1-generator",
            help="Stage 1's generative objective — overrides --mode for stage 1 only",
        ),
    ] = None,
    stage1_hidden_dim: Annotated[
        int | None,
        typer.Option(
            "--stage1-hidden-dim",
            help="Overrides --hidden-dim for stage 1 only (same effect today; "
            "--hidden-dim is kept as a shorthand since stage 1 was the only "
            "target before stage2_model got its own flags)",
        ),
    ] = None,
    stage1_n_res_blocks: Annotated[
        int | None,
        typer.Option("--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"),
    ] = None,
    stage1_dropout: Annotated[
        float | None,
        typer.Option("--stage1-dropout", help="Overrides --dropout for stage 1 only"),
    ] = None,
    stage2_generator: Annotated[
        Mode | None,
        typer.Option(
            "--stage2-generator",
            help="Stage 2's generative objective — overrides --mode for stage 2 "
            "only, e.g. combine with --stage1-generator flow for a mixed "
            "flow/wgan run",
        ),
    ] = None,
    stage2_hidden_dim: Annotated[
        int | None,
        typer.Option("--stage2-hidden-dim", help="Stage 2 trunk width"),
    ] = None,
    stage2_n_res_blocks: Annotated[
        int | None,
        typer.Option("--stage2-n-res-blocks", help="Stage 2 trunk depth"),
    ] = None,
    stage2_dropout: Annotated[
        float | None,
        typer.Option("--stage2-dropout", help="Dropout inside stage 2's ResBlocks"),
    ] = None,
    stage2_decoder: Annotated[
        Decoder | None,
        typer.Option(
            "--stage2-decoder",
            help="one_shot: predict all k_max secondary slots at once (v0.2 "
            "behaviour). autoregressive: emit one secondary at a time in "
            "descending-energy order (default)",
        ),
    ] = None,
    stage2_k_max: Annotated[
        int | None,
        typer.Option(
            "--stage2-k-max",
            help="Maximum secondary slots (fixed width under one_shot, a "
            "generation-loop safety cap under autoregressive; default: 15)",
        ),
    ] = None,
    stage2_context_dim: Annotated[
        int | None,
        typer.Option(
            "--stage2-context-dim",
            help="Width of the projected stage-1 outcome fed into stage 2's conditioning (default: 64)",
        ),
    ] = None,
    stage2_stage1_context: Annotated[
        Stage1Context | None,
        typer.Option(
            "--stage2-stage1-context",
            help="What stage 2 conditions on during training: 'truth' (the "
            "ground-truth stage-1 target, detached — default) or 'sampled' "
            "(stage 1's own sampled output, closing the train/inference gap "
            "at the cost of an extra sampling pass per batch)",
        ),
    ] = None,
    conditioning: Annotated[
        Conditioning | None,
        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[
        bool | None,
        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[
        str | None,
        typer.Option("--router-type", help="Router implementation name (see ROUTER_REGISTRY)"),
    ] = None,
    n_experts: Annotated[int | None, typer.Option("--n-experts", help="Number of routed experts")] = None,
    router_axis: Annotated[
        list[str] | None,
        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[
        int | None,
        typer.Option(
            "--n-critic",
            help="WGAN-GP (--mode wgan only): critic updates per generator update (default: 5)",
        ),
    ] = None,
    gp_weight: Annotated[
        float | None,
        typer.Option(
            "--gp-weight",
            help="WGAN-GP (--mode wgan only): gradient-penalty coefficient (default: 10.0)",
        ),
    ] = None,
    noise_dim: Annotated[
        int | None,
        typer.Option(
            "--noise-dim",
            help="WGAN (--mode wgan only): generator input noise-vector width (default: 64)",
        ),
    ] = None,
    critic_lr: Annotated[
        float | None,
        typer.Option(
            "--critic-lr",
            help="WGAN-GP (--mode wgan only): critic learning rate (default: same as --lr)",
        ),
    ] = None,
    stage1_n_critic: Annotated[
        int | None,
        typer.Option("--stage1-n-critic", help="Overrides --n-critic for stage 1 only"),
    ] = None,
    stage1_gp_weight: Annotated[
        float | None,
        typer.Option("--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"),
    ] = None,
    stage1_noise_dim: Annotated[
        int | None,
        typer.Option("--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"),
    ] = None,
    stage1_critic_lr: Annotated[
        float | None,
        typer.Option("--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"),
    ] = None,
    stage2_n_critic: Annotated[
        int | None,
        typer.Option("--stage2-n-critic", help="Overrides --n-critic for stage 2 only"),
    ] = None,
    stage2_gp_weight: Annotated[
        float | None,
        typer.Option("--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"),
    ] = None,
    stage2_noise_dim: Annotated[
        int | None,
        typer.Option(
            "--stage2-noise-dim",
            help="Overrides --noise-dim for stage 2 only; under "
            "--stage2-decoder autoregressive a fresh draw is made per token",
        ),
    ] = None,
    stage2_critic_lr: Annotated[
        float | None,
        typer.Option("--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"),
    ] = None,
    stage1_critic_hidden_dim: Annotated[
        int | None,
        typer.Option(
            "--stage1-critic-hidden-dim",
            help="WGAN-GP (--mode wgan only): critic width for stage 1 (default: same as generator's hidden_dim)",
        ),
    ] = None,
    stage1_critic_n_res_blocks: Annotated[
        int | None,
        typer.Option(
            "--stage1-critic-n-res-blocks",
            help="WGAN-GP (--mode wgan only): critic depth for stage 1 (default: same as generator's n_res_blocks)",
        ),
    ] = None,
    stage2_critic_hidden_dim: Annotated[
        int | None,
        typer.Option(
            "--stage2-critic-hidden-dim",
            help="WGAN-GP (--mode wgan only): critic width for stage 2 (default: same as generator's hidden_dim)",
        ),
    ] = None,
    stage2_critic_n_res_blocks: Annotated[
        int | None,
        typer.Option(
            "--stage2-critic-n-res-blocks",
            help="WGAN-GP (--mode wgan only): critic depth for stage 2 (default: same as generator's n_res_blocks)",
        ),
    ] = None,
    stage1_init_from: Annotated[
        Path | None,
        typer.Option(
            "--stage1-init-from",
            help="Checkpoint .pt to load stage 1's weights from before training starts "
            "(gitea #42) — combine with --stage1-freeze to retrain stage 2 alone "
            "against a fixed, known-good stage 1",
        ),
    ] = None,
    stage1_freeze: Annotated[
        bool | None,
        typer.Option(
            "--stage1-freeze/--no-stage1-freeze",
            help="Never update stage 1's weights (requires --stage1-init-from, or --resume)",
        ),
    ] = None,
    stage2_init_from: Annotated[
        Path | None,
        typer.Option(
            "--stage2-init-from",
            help="Checkpoint .pt to load stage 2's weights from before training starts (gitea #42)",
        ),
    ] = None,
    stage2_freeze: Annotated[
        bool | None,
        typer.Option(
            "--stage2-freeze/--no-stage2-freeze",
            help="Never update stage 2's weights (requires --stage2-init-from, or --resume)",
        ),
    ] = None,
    val_fraction: Annotated[float | None, typer.Option("--val-fraction", "-f")] = None,
    seed: Annotated[
        int | None,
        typer.Option("--seed", "-s", help="Random seed for reproducibility"),
    ] = None,
    validate_every: Annotated[
        int | None,
        typer.Option(
            "--validate-every",
            "-v",
            help="Run marginal+KL validation every N epochs (0 disables)",
        ),
    ] = None,
    validate_steps: Annotated[
        int | None,
        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[
        int | None,
        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,
    cache_setup: Annotated[
        bool,
        typer.Option(
            "--cache-setup/--no-cache-setup",
            help="Cache the training setup stage's expensive per-file "
            "precomputation (vocab maps, event split index, normalizer stats) "
            "in a JSON sidecar next to the data, so a repeat `giant train` "
            "against the same dataset (e.g. a hyperparameter sweep) can skip "
            "re-deriving it",
        ),
    ] = True,
    rebuild_setup_cache: Annotated[
        bool,
        typer.Option(
            "--rebuild-setup-cache/--no-rebuild-setup-cache",
            help="Ignore any existing setup cache sidecar and recompute every "
            "section fresh for this run (still writes the refreshed sections "
            "back to the sidecar for later runs; no effect if --no-cache-setup)",
        ),
    ] = False,
    out: Annotated[
        Path | None,
        typer.Option(
            "--out",
            "-o",
            help="Checkpoint dir (default: timestamped dir from hyperparams, "
            "or the --resume checkpoint's own dir when resuming)",
        ),
    ] = None,
    device: Annotated[
        str | None,
        typer.Option("--device", "-D", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    num_workers: Annotated[int | None, typer.Option("--num-workers", "-j")] = None,
    resume: Annotated[
        Path | None,
        typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"),
    ] = None,
    wandb: Annotated[
        bool | None,
        typer.Option(
            "--wandb/--no-wandb",
            help="Log per-epoch training metrics to Weights & Biases (requires `uv sync --extra wandb`)",
        ),
    ] = None,
    wandb_project: Annotated[
        str | None,
        typer.Option("--wandb-project", help="W&B project name (default: giant)"),
    ] = None,
    wandb_run_name: Annotated[
        str | None,
        typer.Option("--wandb-run-name", help="W&B run name (default: out_dir name)"),
    ] = None,
    wandb_log_every: Annotated[
        int | None,
        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,
    precision: Annotated[
        str | None,
        typer.Option(
            "--precision",
            help="Training-step autocast precision: 'fp32' (default) or "
            "'bf16'. No 'fp16' — see giant.training.amp.resolve_autocast",
        ),
    ] = None,
) -> None:
    """Train the GIANT surrogate model."""
    import torch

    from giant.pipeline import run_train_job

    batch_size_auto = False
    batch_size_value: int | None = 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', got {batch_size!r}",
                    err=True,
                )
                raise typer.Exit(1)

    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
    flag_values: dict[str, object] = {
        "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,
        "wandb": wandb,
        "wandb_project": wandb_project,
        "wandb_run_name": wandb_run_name,
        "wandb_log_every": wandb_log_every,
        "precision": precision,
        "hidden_dim": hidden_dim,
        "n_blocks": n_blocks,
        "dropout": dropout,
        "stage1_hidden_dim": stage1_hidden_dim,
        "stage1_n_res_blocks": stage1_n_res_blocks,
        "stage1_dropout": stage1_dropout,
        "stage2_hidden_dim": stage2_hidden_dim,
        "stage2_n_res_blocks": stage2_n_res_blocks,
        "stage2_dropout": stage2_dropout,
        "stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
        "stage2_k_max": stage2_k_max,
        "stage2_context_dim": stage2_context_dim,
        "stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
        "mode": mode.value if mode is not None else None,
        "stage1_generator": stage1_generator.value if stage1_generator is not None else None,
        "stage2_generator": stage2_generator.value if stage2_generator is not None else None,
        "conditioning": conditioning.value if conditioning is not None else None,
        "emb_dim": emb_dim,
        "router_config": cli_router or None,
        "n_critic": n_critic,
        "gp_weight": gp_weight,
        "noise_dim": noise_dim,
        "critic_lr": critic_lr,
        "stage1_n_critic": stage1_n_critic,
        "stage1_gp_weight": stage1_gp_weight,
        "stage1_noise_dim": stage1_noise_dim,
        "stage1_critic_lr": stage1_critic_lr,
        "stage2_n_critic": stage2_n_critic,
        "stage2_gp_weight": stage2_gp_weight,
        "stage2_noise_dim": stage2_noise_dim,
        "stage2_critic_lr": stage2_critic_lr,
        "stage1_critic_hidden_dim": stage1_critic_hidden_dim,
        "stage1_critic_n_res_blocks": stage1_critic_n_res_blocks,
        "stage2_critic_hidden_dim": stage2_critic_hidden_dim,
        "stage2_critic_n_res_blocks": stage2_critic_n_res_blocks,
        "stage1_init_from": str(stage1_init_from) if stage1_init_from is not None else None,
        "stage1_freeze": stage1_freeze,
        "stage2_init_from": str(stage2_init_from) if stage2_init_from is not None else None,
        "stage2_freeze": stage2_freeze,
    }
    overrides = gconfig.overrides_from_flags(flag_values)

    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
    gconfig.validate_config(cfg, resume=resume is not None)
    t = cfg["train"]

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

    if batch_size_auto:
        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(cfg, training=True, stage="stage1")
        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.training) — hence the suffix loop in
        # resolve_default_out_dir.
        out_dir = gconfig.resolve_default_out_dir(cfg)

    typer.echo(f"device:  {_device}")
    typer.echo(f"out_dir: {out_dir}")
    typer.echo(f"precision: {t['precision']}")

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


@app.command("new-run")
def new_run(
    config: Annotated[
        Path | None,
        typer.Option(
            "--config",
            "-c",
            help="Base TOML to start from (default: built-in defaults)",
        ),
    ] = None,
    mode: Annotated[Mode | None, typer.Option("--mode", "-m")] = None,
    epochs: Annotated[int | None, typer.Option("--epochs", "-e")] = None,
    batch_size: Annotated[int | None, typer.Option("--batch-size", "-b")] = None,
    lr: Annotated[float | None, typer.Option("--lr", "-l")] = None,
    hidden_dim: Annotated[int | None, typer.Option("--hidden-dim", "-H")] = None,
    n_blocks: Annotated[int | None, typer.Option("--n-blocks", "-n")] = None,
    emb_dim: Annotated[int | None, typer.Option("--emb-dim", "-E")] = None,
    dropout: Annotated[float | None, typer.Option("--dropout", "-d")] = None,
    stage1_generator: Annotated[Mode | None, typer.Option("--stage1-generator")] = None,
    stage1_hidden_dim: Annotated[int | None, typer.Option("--stage1-hidden-dim")] = None,
    stage1_n_res_blocks: Annotated[int | None, typer.Option("--stage1-n-res-blocks")] = None,
    stage1_dropout: Annotated[float | None, typer.Option("--stage1-dropout")] = None,
    stage2_generator: Annotated[Mode | None, typer.Option("--stage2-generator")] = None,
    stage2_hidden_dim: Annotated[int | None, typer.Option("--stage2-hidden-dim")] = None,
    stage2_n_res_blocks: Annotated[int | None, typer.Option("--stage2-n-res-blocks")] = None,
    stage2_dropout: Annotated[float | None, typer.Option("--stage2-dropout")] = None,
    stage2_decoder: Annotated[Decoder | None, typer.Option("--stage2-decoder")] = None,
    stage2_k_max: Annotated[int | None, typer.Option("--stage2-k-max")] = None,
    stage2_context_dim: Annotated[int | None, typer.Option("--stage2-context-dim")] = None,
    stage2_stage1_context: Annotated[Stage1Context | None, typer.Option("--stage2-stage1-context")] = None,
    stage1_init_from: Annotated[Path | None, typer.Option("--stage1-init-from")] = None,
    stage1_freeze: Annotated[bool | None, typer.Option("--stage1-freeze/--no-stage1-freeze")] = None,
    stage2_init_from: Annotated[Path | None, typer.Option("--stage2-init-from")] = None,
    stage2_freeze: Annotated[bool | None, typer.Option("--stage2-freeze/--no-stage2-freeze")] = None,
    conditioning: Annotated[Conditioning | None, typer.Option("--conditioning")] = None,
    router: Annotated[bool | None, typer.Option("--router/--no-router")] = None,
    router_type: Annotated[str | None, typer.Option("--router-type")] = None,
    n_experts: Annotated[int | None, typer.Option("--n-experts")] = None,
    router_axis: Annotated[list[str] | None, typer.Option("--router-axis")] = None,
    out: Annotated[
        Path | None,
        typer.Option("--out", "-o", help="Run dir (default: auto from hyperparams)"),
    ] = None,
    comment: Annotated[
        str | None,
        typer.Option("--comment", help="Free-text note recorded in config.toml's meta section"),
    ] = None,
    data: Annotated[
        Path | None,
        typer.Option(
            "--data",
            help="Dataset path to fill in the printed next-step command (not stored in the config)",
        ),
    ] = None,
    force: Annotated[
        bool,
        typer.Option(
            "--force",
            help="Overwrite config.toml even if --out already has checkpoints",
        ),
    ] = False,
    dry_run: Annotated[
        bool,
        typer.Option("--dry-run", help="Print the resolved config without writing anything"),
    ] = False,
) -> None:
    """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir.

    This is the config-file-first counterpart to hand-editing a TOML: start
    from a base --config (or built-in defaults), override a few hyperparams
    inline, and this resolves+writes the full `config.toml` into a fresh (or
    explicit --out) run dir — the same file `giant train --config ...` reads.
    `giant train` itself overwrites this file in place once it actually runs
    (with the full dataset-derived meta section), so this scaffold's meta
    section is just a placeholder recording what was asked for and when.
    """
    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
    flag_values: dict[str, object] = {
        "epochs": epochs,
        "batch_size": batch_size,
        "lr": lr,
        "hidden_dim": hidden_dim,
        "n_blocks": n_blocks,
        "dropout": dropout,
        "stage1_hidden_dim": stage1_hidden_dim,
        "stage1_n_res_blocks": stage1_n_res_blocks,
        "stage1_dropout": stage1_dropout,
        "stage2_hidden_dim": stage2_hidden_dim,
        "stage2_n_res_blocks": stage2_n_res_blocks,
        "stage2_dropout": stage2_dropout,
        "stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
        "stage2_k_max": stage2_k_max,
        "stage2_context_dim": stage2_context_dim,
        "stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
        "stage1_init_from": str(stage1_init_from) if stage1_init_from is not None else None,
        "stage1_freeze": stage1_freeze,
        "stage2_init_from": str(stage2_init_from) if stage2_init_from is not None else None,
        "stage2_freeze": stage2_freeze,
        "mode": mode.value if mode is not None else None,
        "stage1_generator": stage1_generator.value if stage1_generator is not None else None,
        "stage2_generator": stage2_generator.value if stage2_generator is not None else None,
        "conditioning": conditioning.value if conditioning is not None else None,
        "emb_dim": emb_dim,
        "router_config": cli_router or None,
    }
    overrides = gconfig.overrides_from_flags(flag_values)

    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
    gconfig.validate_config(cfg)
    run_dir = (out or gconfig.resolve_default_out_dir(cfg)).resolve()

    if not force:
        existing = [n for n in ("last.pt", "best.pt") if (run_dir / n).exists()]
        if existing:
            typer.echo(
                f"error: {run_dir} already has {', '.join(existing)} — pass "
                "--force to overwrite its config.toml anyway",
                err=True,
            )
            raise typer.Exit(1)

    typer.echo(f"run dir: {run_dir}")

    if dry_run:
        typer.echo("dry-run: not writing anything. Resolved config:")
        for section in ("train", "conditioning", "stage1_model", "stage2_model"):
            typer.echo(f"[{section}]")
            for k, v in cfg[section].items():
                if k == "router":
                    continue
                typer.echo(f"  {k} = {v}")
        return

    meta = {
        # Tags the written config.toml as v0.3-shaped so a later
        # `migrate_config` load (e.g. `giant train --config ...`) treats it
        # as already-migrated instead of misreading it as v0.2 and dropping
        # its stage1_model/stage2_model/conditioning content.
        "config_version": gconfig.CONFIG_VERSION,
        "git_hash": gconfig.git_hash(),
        "created_at": datetime.now(UTC).isoformat(timespec="seconds"),
        "created_by": "giant new-run",
    }
    if comment:
        meta["comment"] = comment

    run_dir.mkdir(parents=True, exist_ok=True)
    gconfig.save_config(cfg, run_dir, meta)
    config_path = run_dir / "config.toml"
    typer.echo(f"wrote {config_path}")

    data_arg = str(data) if data is not None else "<data.parquet>"
    typer.echo("")
    typer.echo("next:")
    typer.echo(f"  giant train {data_arg} --config {config_path} --out {run_dir}")


model_app = typer.Typer(
    no_args_is_help=True,
    help="Inspect a resolved model architecture without training.",
)
app.add_typer(model_app, name="model")


@model_app.command("summary")
def model_summary(
    config: Annotated[
        Path | None,
        typer.Option("--config", "-c", help="TOML config file (default: built-in defaults)"),
    ] = None,
    pdg_vocab: Annotated[
        int,
        typer.Option(
            "--pdg-vocab",
            help="Placeholder PDG vocab size for conditioning.particle.type='embedding' "
            "or a pdg/process router (no dataset attached to derive the real training vocab)",
        ),
    ] = 300,
    mat_vocab: Annotated[
        int,
        typer.Option(
            "--mat-vocab",
            help="Placeholder material vocab size for conditioning.material.type='embedding' "
            "or a process router (default: the number of known materials in giant.materials)",
        ),
    ] = len(MATERIAL_PROPERTIES),
) -> None:
    """Build the resolved model graph from a config with no dataset attached, and
    print per-module parameter counts, trunk widths, which heads exist, and which
    conditioning/stage1_model/stage2_model config keys actually shaped it."""
    from giant.model.summary import render_summary, summarize_model

    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, {})
    try:
        gconfig.validate_config(cfg)
    except ValueError as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(1)

    summary = summarize_model(cfg, pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    typer.echo(render_summary(summary))


@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[
        Path | None,
        typer.Option(
            "--out",
            "-o",
            help="Output parquet path (default: a UUID-named file under /ceph's central "
            "predictions store if --data is under /ceph, else a sibling of --data)",
        ),
    ] = None,
    truth: Annotated[
        bool,
        typer.Option(
            "--truth/--no-truth",
            help="--coord global only: also read and write ground-truth post-step + "
            "secondary columns (true_step_length, true_edep, true_sec_*_list, ...) "
            "alongside the predictions, at the cost of reading full row-groups instead "
            "of conditioning columns only. Ignored for --coord local, which is always "
            "paired. Default: on.",
        ),
    ] = True,
    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[
        str | None,
        typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    comment: Annotated[
        str | None,
        typer.Option(
            "--comment",
            "-m",
            help="Free-text note recorded in the prediction's YAML sidecar",
        ),
    ] = None,
    set_: Annotated[
        list[str] | None,
        typer.Option(
            "--set",
            help="Override a sampling-only model_config key on this checkpoint, "
            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
        ),
    ] = None,
) -> None:
    """Run trained model on a parquet file and save predictions."""
    import time

    import numpy as np
    import pyarrow as pa
    import pyarrow.parquet as pq
    import torch
    from tqdm import tqdm

    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
    from giant.data.loader import event_id_offset, find_parquet_files, iter_cond_chunks, iter_file_chunks
    from giant.data.transforms import (
        build_cond_features,
        build_features,
        energy_simplex_decode,
        inv_local_frame_rotation,
        inv_log_transform,
        reconstruct_post_pos,
    )
    from giant.rollout import decode_secondary_identity
    from giant.sample import resolve_n_sec, sample_stage1, sample_stage2

    _t_setup_start = time.perf_counter()

    batch_size_auto = False
    batch_size_value: int | None = 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 ---
    config_overrides = _parse_set_flags(set_)
    try:
        ctx = load_for_inference(
            checkpoint, _device, "predict", weights=weights.value, config_overrides=config_overrides
        )
    except CheckpointCompatibilityError as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(1)
    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")

    assert ctx.stage1 is not None and ctx.stage2 is not None  # require_stage2=True (default) guarantees this
    model, sec_decoder = ctx.stage1, ctx.stage2
    cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm
    pdg_map, mat_map = ctx.pdg_map, ctx.mat_map
    pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map
    sec_type_topn_map = ctx.sec_type_topn_map
    particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning
    other_policy = ctx.other_policy
    stage1_ddpm_steps = ctx.stage1_ddpm_steps
    stage2_k_max = ctx.k_max

    if batch_size_auto:
        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(ctx.model_config, 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

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

    training_cfg = gconfig.load_checkpoint_config(checkpoint)
    _write_s = 0.0
    _setup_s = time.perf_counter() - _t_setup_start

    # --coord local always needs full row-groups (it's paired against the 9D
    # target); --coord global only needs them when --truth is requested —
    # otherwise the cheaper conditioning-only read is used.
    write_truth = coord == Coord.global_ and truth

    def chunk_iter(path: Path, offset: int):
        if coord == Coord.local or write_truth:
            return iter_file_chunks(path, offset=offset, k_max=stage2_k_max)
        return iter_cond_chunks(path, offset=offset)

    # load_for_inference already guarantees pdg_topn_map/mat_topn_map are not
    # None whenever the matching conditioning axis is "onehot" — the extra
    # `is not None` conjuncts below are redundant at runtime, just narrowing
    # for the type checker.
    cond_pdg_topn = pdg_topn_map.class_map if pdg_topn_map is not None and particle_conditioning == "onehot" else None
    cond_mat_topn = mat_topn_map.class_map if mat_topn_map is not None and material_conditioning == "onehot" else None

    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, _write_s

        if coord == Coord.local:
            feats = build_features(
                piece,
                pdg_map,
                mat_map,
                particle_conditioning=particle_conditioning,
                material_conditioning=material_conditioning,
                pdg_topn_map=cond_pdg_topn,
                mat_topn_map=cond_mat_topn,
                k_max=stage2_k_max,
            )
            cond_cat = feats.cond_cat
            target_raw = feats.target_s1
            cond_cont = cond_norm.transform(feats.cond_cont)
        else:
            cond_cont, cond_cat = build_cond_features(
                piece,
                pdg_map,
                mat_map,
                cond_norm,
                particle_conditioning=particle_conditioning,
                material_conditioning=material_conditioning,
                pdg_topn_map=cond_pdg_topn,
                mat_topn_map=cond_mat_topn,
            )

        cc = torch.from_numpy(cond_cont).float().to(_device)
        ck = torch.from_numpy(cond_cat).long().to(_device)
        stage1_norm, n_sec_pred = sample_stage1(model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps)

        if coord == Coord.global_:
            # A fresh v0.3.0 Stage1Model owns no n_sec_head —
            # sample_stage1 returns n_sec_pred=None then, so ask stage 2.
            n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred)
            sec_cont, sec_type, sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
            # A stop-token decoder resolves n_sec_pred=None above — read the
            # real count back off sec_valid_pred instead (a no-op round trip
            # under every other n_sec.mode, where sec_valid_pred was built
            # FROM n_sec_pred in the first place).
            n_sec_pred_np = sec_valid_pred.sum(dim=-1).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)

            # particle_type.target="physical": sec_pdg_code is a reporting-
            # only nearest-known-PDG label (never fed back into the model —
            # "no snapping at inference"). "onehot"/"embedding": PDG
            # resolution IS the secondary's identity — see
            # decode_secondary_identity's docstring.
            sec_E, sec_dir_world, _, _, sec_pdg_code, _l1_dist = decode_secondary_identity(
                sec_decoder,
                sec_cont,
                sec_type,
                n_sec_pred_np,
                e_sec_pred,
                piece["pre_dir"],
                sec_phys_norm,
                pdg_map,
                sec_type_topn_map,
                other_policy,
                None,
            )
            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)]

            columns = {
                "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,
            }

            if write_truth:
                n_sec_true = piece["n_sec"]
                columns.update(
                    {
                        "true_step_length": piece["step_length"],
                        "true_delta_e": piece["delta_e"],
                        "true_edep": piece["edep"],
                        "true_post_E": piece["post_E"],
                        "true_post_dx": piece["post_dir"][:, 0],
                        "true_post_dy": piece["post_dir"][:, 1],
                        "true_post_dz": piece["post_dir"][:, 2],
                        "true_post_x": piece["post_pos"][:, 0],
                        "true_post_y": piece["post_pos"][:, 1],
                        "true_post_z": piece["post_pos"][:, 2],
                        "true_e_sec": piece["e_sec"],
                        "process": piece["process"],
                    }
                )
                if "sec_E_list" in piece:
                    columns.update(
                        {
                            "true_sec_pdg_list": [
                                piece["sec_pdg_list"][i, :n].tolist() for i, n in enumerate(n_sec_true)
                            ],
                            "true_sec_E_list": [piece["sec_E_list"][i, :n].tolist() for i, n in enumerate(n_sec_true)],
                            "true_sec_dx_list": [
                                piece["sec_dir_list"][i, :n, 0].tolist() for i, n in enumerate(n_sec_true)
                            ],
                            "true_sec_dy_list": [
                                piece["sec_dir_list"][i, :n, 1].tolist() for i, n in enumerate(n_sec_true)
                            ],
                            "true_sec_dz_list": [
                                piece["sec_dir_list"][i, :n, 2].tolist() for i, n in enumerate(n_sec_true)
                            ],
                        }
                    )

            table = pa.table(columns)

        table = table.replace_schema_metadata(
            {
                PREDICT_COORD_METADATA_KEY: coord.value,
                PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
                PREDICT_TRUTH_METADATA_KEY: "1" if (coord == Coord.local or write_truth) else "0",
            }
        )

        _t0 = time.perf_counter()
        if writer is None:
            writer = pq.ParquetWriter(out, table.schema)
        writer.write_table(table)
        _write_s += time.perf_counter() - _t0
        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

    _t_predict_start = time.perf_counter()
    bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
    for i, path in enumerate(files):
        for chunk in chunk_iter(path, offset=event_id_offset(i)):
            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()
    _predict_s = time.perf_counter() - _t_predict_start

    timing = _build_predict_timing(
        setup_s=_setup_s,
        predict_s=_predict_s,
        write_s=_write_s,
        n_rows=total,
        device=str(_device),
        torch_threads=torch.get_num_threads(),
    )

    ref_path = _write_prediction_ref(
        checkpoint,
        pred_uuid,
        out,
        dataset_path,
        comment,
        extra={
            "coord": coord.value,
            "has_truth": coord == Coord.local or write_truth,
            "schema_version": PREDICT_SCHEMA_VERSION,
            "steps": steps,
            "weights": weights.value,
            "device": str(_device),
            "batch_size": bs,
            "batch_size_auto": batch_size_auto,
            "n_input_rows": total_rows,
            "n_files": len(files),
            "n_rows": total,
            "n_skipped_rows": skipped,
            "unknown_pdg_counts": {str(pdg): count for pdg, count in unknown_pdg_counts.items()},
            "timing": timing,
            # Full architecture spec baked into the checkpoint — see the
            # matching comment in `rollout`.
            "model_config": dict(ctx.model_config),
            "config_overrides": dict(ctx.config_overrides),
            "training_epoch": ctx.epoch,
            "best_val_loss": ctx.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", {})),
        },
    )
    typer.echo(f"reference: {ref_path}")
    if timing["us_per_step"] is not None:
        typer.echo(
            f"timing: {_predict_s:.1f}s total ({timing['sample_s']:.1f}s sample + {_write_s:.1f}s write), "
            f"{timing['us_per_step']:.1f} us/step over {total:,} step(s)"
        )

    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.
    """
    import numpy as np

    from giant.data.loader import event_id_offset, iter_cond_chunks

    best_E: dict[int, float] = {}
    best: dict[int, tuple] = {}
    for file_idx, path in enumerate(files):
        for chunk in iter_cond_chunks(path, offset=event_id_offset(file_idx)):
            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[
        int | None,
        typer.Option(
            "--max-tracks-per-event",
            help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
        ),
    ] = None,
    escape_threshold: Annotated[
        float | None,
        typer.Option(
            "--escape-threshold",
            help="Override the oracle's NN-distance escape threshold [mm]",
        ),
    ] = None,
    n_events: Annotated[int | None, typer.Option("--n-events", help="Cap number of seed events")] = None,
    device: Annotated[str | None, typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")] = None,
    out: Annotated[Path | None, typer.Option("--out", "-o", help="Output steps parquet")] = None,
    seed: Annotated[
        int | None,
        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
    ] = None,
    set_: Annotated[
        list[str] | None,
        typer.Option(
            "--set",
            help="Override a sampling-only model_config key on this checkpoint, "
            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
        ),
    ] = None,
) -> None:
    """Roll the surrogate forward into full showers (autoregressive)."""
    import time

    import numpy as np
    import pyarrow as pa
    import pyarrow.parquet as pq
    import torch

    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
    from giant.data.loader import find_parquet_files
    from giant.geometry import GeometryOracle
    from giant.rollout import L1DistCollector, RolloutSummary
    from giant.rollout import rollout as run_rollout

    _t_setup_start = time.perf_counter()

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

    config_overrides = _parse_set_flags(set_)
    try:
        ctx = load_for_inference(
            checkpoint, _device, "rollout", weights=weights.value, config_overrides=config_overrides
        )
    except CheckpointCompatibilityError as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(1)
    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")

    training_cfg = gconfig.load_checkpoint_config(checkpoint)

    assert ctx.stage1 is not None and ctx.stage2 is not None  # require_stage2=True (default) guarantees this
    model, sec_decoder = ctx.stage1, ctx.stage2
    cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm
    pdg_map, mat_map = ctx.pdg_map, ctx.mat_map
    pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map
    sec_type_topn_map = ctx.sec_type_topn_map
    particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning
    other_policy = ctx.other_policy
    stage1_ddpm_steps, stage2_ddpm_steps = ctx.stage1_ddpm_steps, ctx.stage2_ddpm_steps
    model_cfg = ctx.model_config

    oracle = GeometryOracle.load(geometry)
    typer.echo(f"loaded geometry oracle: {geometry} (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
    _write_s = 0.0

    def _write_chunk(row: dict[str, np.ndarray]) -> None:
        nonlocal writer, _write_s
        _t0 = time.perf_counter()
        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)
        _write_s += time.perf_counter() - _t0

    # Only meaningful under particle_type.target="embedding" — a
    # no-op collector otherwise, cheaper than branching the call itself.
    l1_dist_collector = L1DistCollector()

    _setup_s = time.perf_counter() - _t_setup_start
    _t_rollout_start = time.perf_counter()
    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,
        particle_conditioning=particle_conditioning,
        material_conditioning=material_conditioning,
        pdg_topn_map=pdg_topn_map,
        mat_topn_map=mat_topn_map,
        sec_type_topn_map=sec_type_topn_map,
        other_policy=other_policy,
        seed=seed,
        stage1_ddpm_steps=stage1_ddpm_steps,
        stage2_ddpm_steps=stage2_ddpm_steps,
        l1_dist_collector=l1_dist_collector,
    )
    if writer is not None:
        writer.close()
    # on_chunk=_write_chunk is always passed above, so rollout() always
    # returns the streaming-summary shape (RolloutSummary), never the
    # materialized dict[str, np.ndarray] alternative its return type allows.
    summary = cast(RolloutSummary, summary)
    _rollout_s = time.perf_counter() - _t_rollout_start
    timing = _build_rollout_timing(
        setup_s=_setup_s,
        rollout_s=_rollout_s,
        write_s=_write_s,
        n_rows=summary["n_rows"],
        termination_reason_counts=summary["termination_reason_counts"],
        n_seed_events=len(seeds["event_id"]),
        device=str(_device),
        torch_threads=torch.get_num_threads(),
    )
    _sample_s = timing["sample_s"]
    n_physical_rows = timing["n_physical_rows"]

    l1_summary = l1_dist_collector.summary()

    ref_path = _write_prediction_ref(
        checkpoint,
        pred_uuid,
        out,
        dataset_path,
        extra={
            "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": 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"],
            # Wall-clock cost of this run, normalized per physical step (the
            # comparable unit against giant.analysis.geant4_reference) — see
            # eval_cost_per_step in giant/analysis/catalog.py.
            "timing": timing,
            # Diagnostic — only present under
            # stage2_model.particle_type.target="embedding"; omitted (not
            # written as null) otherwise, so giant.analysis can tell "not
            # applicable to this checkpoint" apart from "collector empty".
            **({"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}),
            # 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),
            "config_overrides": dict(ctx.config_overrides),
            "training_epoch": ctx.epoch,
            "best_val_loss": ctx.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", {})),
        },
    )

    typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
    typer.echo(f"terminations: {summary['termination_reason_counts']}")
    if timing["us_per_step"] is not None:
        typer.echo(
            f"timing: {_rollout_s:.1f}s total ({_sample_s:.1f}s sample + {_write_s:.1f}s write), "
            f"{timing['us_per_step']:.1f} us/step over {n_physical_rows:,} physical steps"
        )
    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_yamls: Annotated[
        list[Path],
        typer.Argument(
            help="giant rollout YAML sidecar(s) (names the rollout + reference files). "
            "Multiple compare N rollouts against one shared reference — every YAML must "
            "name the same `dataset`."
        ),
    ],
    label: Annotated[
        list[str] | None,
        typer.Option(
            "--label",
            help="Series name for a rollout YAML, positionally matched to it — give none, "
            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
            "single YAML).",
        ),
    ] = None,
    prediction: Annotated[
        list[Path] | None,
        typer.Option(
            "--prediction",
            help="giant predict YAML sidecar(s) (paired truth/pred comparison, the "
            "`prediction` plot family) — optional add-on to the rollout comparison. "
            "Every one must be seeded from the same `dataset` as the rollout(s) and "
            "share one predict --coord.",
        ),
    ] = None,
    prediction_label: Annotated[
        list[str] | None,
        typer.Option(
            "--prediction-label",
            help="Series name for a --prediction YAML, positionally matched to it — give "
            'none, or exactly one per YAML. Defaults to the YAML stem (or "prediction" '
            "for a single YAML).",
        ),
    ] = None,
    run_dir: Annotated[
        Path | None,
        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 (+ optional prediction) YAML(s) → shared.json + run_meta.json."""
    from giant.analysis import prep

    path = prep(
        rollout_yamls,
        run_dir,
        n_chunks=chunks,
        default_base=Path.cwd() / "analysis_runs",
        labels=label,
        prediction_yamls=prediction or (),
        prediction_labels=prediction_label,
        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("metrics")
def analyze_metrics(
    run_dir: Annotated[Path, typer.Argument(help="Run directory containing metrics.csv (from `giant train`)")],
    out_dir: Annotated[
        Path | None,
        typer.Option(
            "--out",
            "-o",
            help="Override the output directory (default: <cwd>/analysis_runs/metrics_<run_dir name>)",
        ),
    ] = None,
) -> None:
    """Render training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) from <run_dir>/metrics.csv."""
    from giant.training.plots import render_metrics

    paths = render_metrics(run_dir, out_dir, default_base=Path.cwd() / "analysis_runs")
    typer.echo(f"rendered {len(paths)} plots -> {paths[0].parent if paths else '(nothing to render)'}")


@analyze_app.command("submit")
def analyze_submit(
    rollout_yamls: Annotated[
        list[Path],
        typer.Argument(
            help="giant rollout YAML sidecar(s). Multiple compare N rollouts against one "
            "shared reference — every YAML must name the same `dataset`."
        ),
    ],
    accounting_group: Annotated[str, typer.Option("--accounting-group")],
    label: Annotated[
        list[str] | None,
        typer.Option(
            "--label",
            help="Series name for a rollout YAML, positionally matched to it — give none, "
            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
            "single YAML).",
        ),
    ] = None,
    prediction: Annotated[
        list[Path] | None,
        typer.Option(
            "--prediction",
            help="giant predict YAML sidecar(s) (paired truth/pred comparison, the "
            "`prediction` plot family) — optional add-on to the rollout comparison. "
            "Every one must be seeded from the same `dataset` as the rollout(s) and "
            "share one predict --coord.",
        ),
    ] = None,
    prediction_label: Annotated[
        list[str] | None,
        typer.Option(
            "--prediction-label",
            help="Series name for a --prediction YAML, positionally matched to it — give "
            'none, or exactly one per YAML. Defaults to the YAML stem (or "prediction" '
            "for a single YAML).",
        ),
    ] = None,
    run_dir: Annotated[
        Path | None,
        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_yamls,
        run_dir,
        n_chunks=chunks,
        default_base=Path.cwd() / "analysis_runs",
        labels=label,
        prediction_yamls=prediction or (),
        prediction_labels=prediction_label,
        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()