giant/giant/config.py

from __future__ import annotations

import copy
import difflib
import hashlib
import random
import re
import subprocess
import sys
import tomllib
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING

from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing

if TYPE_CHECKING:
    import torch


class Conditioning(str, Enum):
    """`conditioning.particle.type` / `conditioning.material.type` choices —
    shared by `giant.cli` and `giant.tools.dwarf`'s Typer commands so the two
    CLIs can't silently drift apart on the option's valid values (see
    DEFAULT_CONFIG["conditioning"] for what each value means)."""

    physical = "physical"
    embedding = "embedding"
    onehot = "onehot"


# Tags a config dict (config.toml, or a checkpoint's model_config) as the new
# v0.3 nested format. Absence of `[meta].config_version == CONFIG_VERSION` is
# read as "this is a v0.2 dict" by migrate_config below.
CONFIG_VERSION = 3


# --- Config dataclasses -----------------------------------------------------
#
# These are the single source of truth for every default below. DEFAULT_CONFIG
# (a plain dict, for merge_cli_overrides/save_config/TOML round-tripping) is
# *generated* from GiantConfig().to_dict() rather than hand-maintained, so it
# cannot drift from the fallback defaults that build_models/build_critics
# (giant/model/network.py) and StageSpec.from_config (giant/training/trainers.py)
# read off these same dataclasses — see issues.md Issue 1.
#
# Each dataclass is frozen and carries an explicit from_dict/to_dict pair
# (mirroring StageSpec's established style in trainers.py) rather than a
# generic reflection-based helper, so every default is fully type-checkable.
# `lambda` is a Python keyword, so dict key "lambda" is always exposed as the
# field `lambda_weight`.
#
# The router sub-block carries genuinely dynamic keys that don't fit a fixed schema:
# composed-router `axis{i}_{field}` flags (see giant.model.network._parse_composed_axes)
# and pipeline.py's runtime-seeded `centers_init`. It carries an `extra: dict` catch-all
# so these keys round-trip losslessly without becoming named fields that would leak into
# every new run's config.toml.


@dataclass(frozen=True)
class ConditioningAxisConfig:
    """One conditioning axis: `conditioning.particle` or `conditioning.material`."""

    # "physical": a small MLP over log(mass)/charge, computable for any PDG
    # code — generalizes beyond the training menu.
    # "embedding": a learned nn.Embedding over a dense training-vocab index —
    # memorizes the training menu; the generalization-comparison baseline,
    # and the only mode compatible with stage2_model.particle_type.target =
    # "embedding".
    # "onehot": a fixed, unlearned vector — top (emb_dim - 1) PDG codes by
    # training-set count, plus one "other" bin. NOT a reparameterization of
    # "embedding": the vocabulary cap is the real difference.
    type: str = "physical"
    # Width of this axis's vector. Under "onehot" this also sets the class
    # count.
    emb_dim: int = 16
    # Depth of the sub-MLP under "physical". Ignored under "embedding"/"onehot".
    n_layers: int = 1

    @classmethod
    def from_dict(cls, d: dict | None) -> "ConditioningAxisConfig":
        d = d or {}
        return cls(
            type=d.get("type", "physical"),
            emb_dim=d.get("emb_dim", 16),
            n_layers=d.get("n_layers", 1),
        )

    def to_dict(self) -> dict:
        return {"type": self.type, "emb_dim": self.emb_dim, "n_layers": self.n_layers}


@dataclass(frozen=True)
class ConditioningConfig:
    # Width of the fused conditioning vector produced by the encoder's fusion
    # MLP, consumed by every downstream trunk.
    out_dim: int = 128
    # false: stage 1 and stage 2 each construct their own ConditionEncoder
    # with identical config but independent weights. true: one instance,
    # shared by reference (halves the conditioning parameter count, forces a
    # common representation).
    share_stages: bool = False
    particle: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig)
    material: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "ConditioningConfig":
        d = d or {}
        return cls(
            out_dim=d.get("out_dim", 128),
            share_stages=d.get("share_stages", False),
            particle=ConditioningAxisConfig.from_dict(d.get("particle")),
            material=ConditioningAxisConfig.from_dict(d.get("material")),
        )

    def to_dict(self) -> dict:
        return {
            "out_dim": self.out_dim,
            "share_stages": self.share_stages,
            "particle": self.particle.to_dict(),
            "material": self.material.to_dict(),
        }


@dataclass(frozen=True)
class FlowConfig:
    # Width of the SinusoidalEmbedding for the flow time variable.
    time_dim: int = 64

    @classmethod
    def from_dict(cls, d: dict | None) -> "FlowConfig":
        d = d or {}
        return cls(time_dim=d.get("time_dim", 64))

    def to_dict(self) -> dict:
        return {"time_dim": self.time_dim}


@dataclass(frozen=True)
class DdpmConfig:
    time_dim: int = 64
    n_steps: int = 1000

    @classmethod
    def from_dict(cls, d: dict | None) -> "DdpmConfig":
        d = d or {}
        return cls(time_dim=d.get("time_dim", 64), n_steps=d.get("n_steps", 1000))

    def to_dict(self) -> dict:
        return {"time_dim": self.time_dim, "n_steps": self.n_steps}


@dataclass(frozen=True)
class Stage1WganConfig:
    noise_dim: int = 64
    n_critic: int = 5
    gp_weight: float = 10.0
    # 0.0 means "inherit train.lr" — not None, since the TOML writer has no
    # null literal to round-trip.
    critic_lr: float = 0.0
    # 0 means "inherit stage1_model.hidden_dim/n_res_blocks" — same
    # round-trip-friendly sentinel as critic_lr above.
    critic_hidden_dim: int = 0
    critic_n_res_blocks: int = 0

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage1WganConfig":
        d = d or {}
        return cls(
            noise_dim=d.get("noise_dim", 64),
            n_critic=d.get("n_critic", 5),
            gp_weight=d.get("gp_weight", 10.0),
            critic_lr=d.get("critic_lr", 0.0),
            critic_hidden_dim=d.get("critic_hidden_dim", 0),
            critic_n_res_blocks=d.get("critic_n_res_blocks", 0),
        )

    def to_dict(self) -> dict:
        return {
            "noise_dim": self.noise_dim,
            "n_critic": self.n_critic,
            "gp_weight": self.gp_weight,
            "critic_lr": self.critic_lr,
            "critic_hidden_dim": self.critic_hidden_dim,
            "critic_n_res_blocks": self.critic_n_res_blocks,
        }


@dataclass(frozen=True)
class Stage2WganConfig(Stage1WganConfig):
    # Straight-through Gumbel temperature for the particle-type one-hot
    # (distinct from router.gumbel_tau_start/_end, which anneal
    # expert-combination weights). Read only under particle_type.target =
    # "onehot".
    gumbel_tau_start: float = 1.0
    gumbel_tau_end: float = 0.1

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage2WganConfig":
        d = d or {}
        return cls(
            noise_dim=d.get("noise_dim", 64),
            n_critic=d.get("n_critic", 5),
            gp_weight=d.get("gp_weight", 10.0),
            critic_lr=d.get("critic_lr", 0.0),
            critic_hidden_dim=d.get("critic_hidden_dim", 0),
            critic_n_res_blocks=d.get("critic_n_res_blocks", 0),
            gumbel_tau_start=d.get("gumbel_tau_start", 1.0),
            gumbel_tau_end=d.get("gumbel_tau_end", 0.1),
        )

    def to_dict(self) -> dict:
        return {
            **super().to_dict(),
            "gumbel_tau_start": self.gumbel_tau_start,
            "gumbel_tau_end": self.gumbel_tau_end,
        }


# Fixed router fields shared by stage1_model.router and stage2_model.router.
# Composed-router axis{i}_{field} keys and pipeline.py's runtime-seeded
# centers_init are NOT in this set — they land in RouterConfig.extra instead
# (see giant.model.network._parse_composed_axes).
_ROUTER_KNOWN_KEYS = frozenset(
    {
        "enabled",
        "type",
        "n_experts",
        "temperature",
        "learn_centers",
        "learn_width",
        "learn_temperature",
        "width_min_ratio",
        "width_max_ratio",
        "lambda_balance",
        "lambda_entropy",
        "gumbel",
        "gumbel_tau_start",
        "gumbel_tau_end",
        "emb_dim",
        "hidden_dim",
        "lambda_proc",
    }
)


@dataclass(frozen=True)
class RouterConfig:
    """`stage1_model.router`'s fixed fields. `extra` holds any key not named
    below — composed-router `axis{i}_{field}` flags and pipeline.py's
    runtime-seeded `centers_init` — so from_dict()/to_dict() round-trip
    losslessly without this dataclass needing to know about them."""

    enabled: bool = False
    type: str = "energy"  # selects the Router impl from ROUTER_REGISTRY
    n_experts: int = 4
    temperature: float = 0.5  # energy/pdg-router kwarg
    learn_centers: bool = True  # energy/pdg-router kwarg
    # energy-router kwargs: mutually exclusive optional learnable
    # gate-sharpness modes. learn_width generalizes the shared `temperature`
    # to one learnable width per expert; learn_temperature instead makes the
    # single shared `temperature` itself learnable. Both are bounded to
    # [width_min_ratio, width_max_ratio] * temperature so gate sharpness
    # can't run away to a collapse-inducing extreme during training.
    learn_width: bool = False
    learn_temperature: bool = False
    width_min_ratio: float = 0.1
    width_max_ratio: float = 10.0
    # Importance-CV^2 load-balancing aux loss weight (Shazeer et al. 2017).
    lambda_balance: float = 0.0
    # Entropy-regularization weight penalizing uniform/collapsed gating — a
    # secondary guard against all experts' widths co-inflating together,
    # which lambda_balance alone can't see.
    lambda_entropy: float = 0.0
    # Opt-in straight-through Gumbel-softmax train-time combine weights: the
    # training forward pass samples a hard one-hot combination (matching
    # eval-time top-1 dispatch exactly) while the backward pass still flows
    # smooth gradient to every expert.
    gumbel: bool = False
    gumbel_tau_start: float = 1.0
    gumbel_tau_end: float = 0.1
    emb_dim: int = 8  # process/pdg-router kwarg: own pdg(/mat) embedding width
    hidden_dim: int = 64  # process-router kwarg: its classifier's hidden width
    lambda_proc: float = 0.0  # process-router kwarg: supervised process-CE weight
    # type = "composed" routes on multiple axes at once (e.g. energy x pdg),
    # each with its own expert count/hyperparameters. Axes are NOT in these
    # defaults — set them as flat axis{i}_{field} keys instead of
    # "n_experts", e.g. axis0_type = "energy", axis0_n_experts = 4,
    # axis1_type = "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See
    # giant.model.network._parse_composed_axes.
    extra: dict = field(default_factory=dict)

    @classmethod
    def from_dict(cls, d: dict | None) -> "RouterConfig":
        d = d or {}
        return cls(
            enabled=d.get("enabled", False),
            type=d.get("type", "energy"),
            n_experts=d.get("n_experts", 4),
            temperature=d.get("temperature", 0.5),
            learn_centers=d.get("learn_centers", True),
            learn_width=d.get("learn_width", False),
            learn_temperature=d.get("learn_temperature", False),
            width_min_ratio=d.get("width_min_ratio", 0.1),
            width_max_ratio=d.get("width_max_ratio", 10.0),
            lambda_balance=d.get("lambda_balance", 0.0),
            lambda_entropy=d.get("lambda_entropy", 0.0),
            gumbel=d.get("gumbel", False),
            gumbel_tau_start=d.get("gumbel_tau_start", 1.0),
            gumbel_tau_end=d.get("gumbel_tau_end", 0.1),
            emb_dim=d.get("emb_dim", 8),
            hidden_dim=d.get("hidden_dim", 64),
            lambda_proc=d.get("lambda_proc", 0.0),
            extra={k: v for k, v in d.items() if k not in _ROUTER_KNOWN_KEYS},
        )

    def to_dict(self) -> dict:
        return {
            "enabled": self.enabled,
            "type": self.type,
            "n_experts": self.n_experts,
            "temperature": self.temperature,
            "learn_centers": self.learn_centers,
            "learn_width": self.learn_width,
            "learn_temperature": self.learn_temperature,
            "width_min_ratio": self.width_min_ratio,
            "width_max_ratio": self.width_max_ratio,
            "lambda_balance": self.lambda_balance,
            "lambda_entropy": self.lambda_entropy,
            "gumbel": self.gumbel,
            "gumbel_tau_start": self.gumbel_tau_start,
            "gumbel_tau_end": self.gumbel_tau_end,
            "emb_dim": self.emb_dim,
            "hidden_dim": self.hidden_dim,
            "lambda_proc": self.lambda_proc,
            **self.extra,
        }


@dataclass(frozen=True)
class TrunkConfig:
    """`stage1_model.trunk`/`stage2_model.trunk`: selects the trunk's expert
    *body* architecture from `giant.model.trunks.TRUNK_REGISTRY` (default
    `"resmlp"` — today's only body, `input_proj -> ResBlock stack ->
    out_proj`). Orthogonal to whether that body is mixed: mixing is still
    controlled entirely by `router.enabled`/`router.n_experts` on the same
    stage, unaffected by this block. A future body's own hyperparameters
    (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling
    field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by
    `generator`.

    `block_conditioning` selects each body's conditioning-injection mechanism
    from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's
    conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour),
    `"film"`, or `"adaln"`."""

    type: str = "resmlp"
    block_conditioning: str = "add"

    @classmethod
    def from_dict(cls, d: dict | None) -> "TrunkConfig":
        d = d or {}
        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))

    def to_dict(self) -> dict:
        return {"type": self.type, "block_conditioning": self.block_conditioning}


@dataclass(frozen=True)
class Stage2RouterConfig(RouterConfig):
    # true: stage 2 shares stage 1's Router module instance, so expert i in
    # stage 1 and expert i in stage 2 gate on identical conditions by
    # construction — every other key in this block is then ignored. false:
    # an independent router.
    tie_to_stage1: bool = False

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage2RouterConfig":
        d = d or {}
        known = _ROUTER_KNOWN_KEYS | {"tie_to_stage1"}
        return cls(
            tie_to_stage1=d.get("tie_to_stage1", False),
            enabled=d.get("enabled", False),
            type=d.get("type", "energy"),
            n_experts=d.get("n_experts", 4),
            temperature=d.get("temperature", 0.5),
            learn_centers=d.get("learn_centers", True),
            learn_width=d.get("learn_width", False),
            learn_temperature=d.get("learn_temperature", False),
            width_min_ratio=d.get("width_min_ratio", 0.1),
            width_max_ratio=d.get("width_max_ratio", 10.0),
            lambda_balance=d.get("lambda_balance", 0.0),
            lambda_entropy=d.get("lambda_entropy", 0.0),
            gumbel=d.get("gumbel", False),
            gumbel_tau_start=d.get("gumbel_tau_start", 1.0),
            gumbel_tau_end=d.get("gumbel_tau_end", 0.1),
            emb_dim=d.get("emb_dim", 8),
            hidden_dim=d.get("hidden_dim", 64),
            lambda_proc=d.get("lambda_proc", 0.0),
            extra={k: v for k, v in d.items() if k not in known},
        )

    def to_dict(self) -> dict:
        return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()}


# stage2_model.n_sec.sampling choices — single source of truth for both
# validate_config's train-time check and INFERENCE_OVERRIDES below.
STOP_SAMPLING_CHOICES = ("greedy", "sample")

# stage2_model.particle_type.other_policy choices — see ParticleTypeConfig's
# docstring for what each means; only documented there until now, since
# nothing validated it at train time.
OTHER_POLICY_CHOICES = ("sample", "modal", "drop")


@dataclass(frozen=True)
class NSecConfig:
    # "head": a classifier over {0..k_max} on the condition encoding alone
    # (no diffusion noise), callable independently at inference.
    # "stop_token": an EOS-style per-slot stop head on the autoregressive
    # secondary decoder (Stage2Autoregressive only — see validate_config),
    # evaluated against the generated prefix instead of conditioning alone.
    # Replaces n_sec_head entirely: the two are mutually exclusive, so this
    # mode builds no n_sec_head and stage2_model.n_sec.lambda instead weights
    # the stop head's BCE term.
    # "truth": take n_sec from ground truth — standalone stage-2 evaluation
    # only, never for rollout.
    mode: str = "head"
    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy/BCE weight for the head
    # Which stage's module physically owns the n_sec_head weights: "stage2" (default,
    # fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1"
    # (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose
    # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
    # to stay attached there, not just be labeled as such).
    owner: str = "stage2"
    # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
    # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
    # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
    # from softmax(logits) (a real sample from the learned count distribution). mode=
    # "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
    # "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
    # "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
    # appears in existing checkpoints' model_config.
    sampling: str = "greedy"

    @classmethod
    def from_dict(cls, d: dict | None) -> "NSecConfig":
        d = d or {}
        return cls(
            mode=d.get("mode", "head"),
            lambda_weight=d.get("lambda", 0.1),
            owner=d.get("owner", "stage2"),
            sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
        )

    def to_dict(self) -> dict:
        return {
            "mode": self.mode,
            "lambda": self.lambda_weight,
            "owner": self.owner,
            "sampling": self.sampling,
        }


@dataclass(frozen=True)
class ParticleTypeConfig:
    # The three targets mirror the three conditioning.particle modes.
    # "onehot": class logits over conditioning.particle.emb_dim classes.
    # "physical": regressed (log mass, charge). "embedding": regressed
    # against conditioning's own particle embedding table (requires
    # conditioning.particle.type = "embedding").
    target: str = "onehot"
    lambda_weight: float = 1.0  # dict key "lambda"
    # How a predicted "other" class becomes a concrete PDG code at rollout.
    # "sample": draw from the empirical within-bucket distribution recorded
    # at map-build time. "modal": always the most common member. "drop":
    # discard the secondary. Read only under target = "onehot".
    other_policy: str = "sample"
    # Secondary-species class count under target = "onehot" — independent of
    # conditioning.particle.emb_dim (see gitea #29: the two used to be
    # silently the same number). 0 = inherit conditioning.particle.emb_dim,
    # preserving pre-#29 behavior.
    n_classes: int = 0
    # Class-balances the target = "onehot" cross-entropy loss against the
    # secondary-species long tail (gitea #44: the failure mode motivating the
    # v0.3.0 pivot was specifically a species collapse — zero photon
    # secondaries, hallucinated antineutrinos). "none": plain CE (pre-#44
    # behavior). "inverse_freq": CE weighted by 1/count per class,
    # normalized to mean 1 across classes so lambda_weight doesn't need
    # retuning when this is switched on. validate_config requires target =
    # "onehot" and stage2_model.generator != "wgan" whenever this isn't
    # "none" — "embedding"/"physical" have no class CE to weight, and the
    # WGAN stage-2 path feeds its type slice to the critic via a
    # straight-through Gumbel relaxation instead of a CE loss.
    class_weighting: str = "none"

    @classmethod
    def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
        d = d or {}
        return cls(
            target=d.get("target", "onehot"),
            lambda_weight=d.get("lambda", 1.0),
            other_policy=d.get("other_policy", "sample"),
            n_classes=d.get("n_classes", 0),
            class_weighting=d.get("class_weighting", "none"),
        )

    def to_dict(self) -> dict:
        return {
            "target": self.target,
            "lambda": self.lambda_weight,
            "other_policy": self.other_policy,
            "n_classes": self.n_classes,
            "class_weighting": self.class_weighting,
        }


@dataclass(frozen=True)
class AutoregressiveConfig:
    # Canonical generation order. Single-valued for now; the key exists so
    # an alternative ordering is not a config break.
    order: str = "energy_desc"
    # How token i+1 sees tokens <= i. "markov": previous token plus running
    # scalars (remaining energy budget, slot index) — a fixed-width summary.
    # "attention": causal self-attention over all emitted tokens.
    history: str = "markov"
    # "always": condition on the ground-truth previous secondary throughout
    # training. "scheduled": scheduled sampling — interpolate toward the
    # model's own prediction. "never": free-running from the start.
    teacher_forcing: str = "always"
    tf_p_start: float = 1.0
    tf_p_end: float = 1.0
    attn_n_heads: int = 4
    attn_n_layers: int = 2

    @classmethod
    def from_dict(cls, d: dict | None) -> "AutoregressiveConfig":
        d = d or {}
        return cls(
            order=d.get("order", "energy_desc"),
            history=d.get("history", "markov"),
            teacher_forcing=d.get("teacher_forcing", "always"),
            tf_p_start=d.get("tf_p_start", 1.0),
            tf_p_end=d.get("tf_p_end", 1.0),
            attn_n_heads=d.get("attn_n_heads", 4),
            attn_n_layers=d.get("attn_n_layers", 2),
        )

    def to_dict(self) -> dict:
        return {
            "order": self.order,
            "history": self.history,
            "teacher_forcing": self.teacher_forcing,
            "tf_p_start": self.tf_p_start,
            "tf_p_end": self.tf_p_end,
            "attn_n_heads": self.attn_n_heads,
            "attn_n_layers": self.attn_n_layers,
        }


@dataclass(frozen=True)
class HeadConfig:
    """A single classifier head's shape — `n_sec_head`/`type_head` (gitea
    #36 deduplicated their five identical hand-rolled
    `Linear -> SiLU -> Linear` definitions into
    `giant.model.layers.build_mlp_head`, which this config drives).
    `hidden_ratio=0.5`/`depth=2` are the exact pre-#36 hardcoded values
    (hidden width = `hidden_dim // 2`, one hidden layer), so omitting a
    `heads` block — including every migrated v0.2 config — reproduces the
    old architecture bit-for-bit."""

    hidden_ratio: float = 0.5  # hidden width = round(hidden_dim * hidden_ratio)
    depth: int = 2  # matches build_mlp_head's depth

    @classmethod
    def from_dict(cls, d: dict | None) -> "HeadConfig":
        d = d or {}
        return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))

    def to_dict(self) -> dict:
        return {"hidden_ratio": self.hidden_ratio, "depth": self.depth}


@dataclass(frozen=True)
class Stage1HeadsConfig:
    """Stage 1 only ever owns `n_sec_head`, and only for a migrated v0.2
    checkpoint (`stage2_model.n_sec.owner = "stage1"`) — see
    `Stage1Model`'s docstring."""

    n_sec: HeadConfig = field(default_factory=HeadConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
        d = d or {}
        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))

    def to_dict(self) -> dict:
        return {"n_sec": self.n_sec.to_dict()}


@dataclass(frozen=True)
class Stage2HeadsConfig:
    """`n_sec` and `type` are independently configurable — n_sec accuracy
    and secondary-species accuracy are separately known weak spots (gitea
    #36)."""

    n_sec: HeadConfig = field(default_factory=HeadConfig)
    type: HeadConfig = field(default_factory=HeadConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
        d = d or {}
        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))

    def to_dict(self) -> dict:
        return {"n_sec": self.n_sec.to_dict(), "type": self.type.to_dict()}


@dataclass(frozen=True)
class Stage1ModelConfig:
    # false skips building/training stage 1 entirely. The resulting
    # checkpoint holds only stage 2 and cannot be rolled out.
    active: bool = True
    # Checkpoint .pt to load this stage's weights from before training starts
    # (its own "model"/"sec_decoder" key, not this run's own resume state) —
    # "" means start from a fresh init. See `freeze` below for the partial-
    # retrain use case this exists for (gitea #42).
    init_from: str = ""
    # true keeps this stage's weights exactly as loaded from `init_from` —
    # forward/backward still run every batch (so its loss/grad_norm metrics
    # stay meaningful, and a WGAN stage's critic still gets a real signal to
    # report), but its optimizer never steps. Lets a rollout-capable
    # checkpoint retrain only the *other* stage against a fixed, known-good
    # one (gitea #42) — `validate_config` requires `init_from` to be set
    # whenever this is true, unless the run is a `--resume`.
    freeze: bool = False
    # "flow": conditional flow matching (~10 ODE steps at inference).
    # "ddpm": cosine-schedule diffusion baseline.
    # "wgan": WGAN-GP, single forward pass at inference.
    generator: str = "flow"
    # Trunk width — also the width of every expert under a routed trunk.
    hidden_dim: int = 256
    # Number of ResBlocks in the trunk, and in every expert under a routed
    # trunk.
    n_res_blocks: int = 6
    dropout: float = 0.0
    # Weight of this stage's loss in the total when both stages are active
    # and non-adversarial. A WGAN stage's adversarial loss drives its own
    # optimizer, so `lambda` scales only its non-adversarial auxiliary terms.
    lambda_weight: float = 1.0  # dict key "lambda"
    flow: FlowConfig = field(default_factory=FlowConfig)
    ddpm: DdpmConfig = field(default_factory=DdpmConfig)
    wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig)
    router: RouterConfig = field(default_factory=RouterConfig)
    trunk: TrunkConfig = field(default_factory=TrunkConfig)
    heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
        d = d or {}
        return cls(
            active=d.get("active", True),
            init_from=d.get("init_from", ""),
            freeze=d.get("freeze", False),
            generator=d.get("generator", "flow"),
            hidden_dim=d.get("hidden_dim", 256),
            n_res_blocks=d.get("n_res_blocks", 6),
            dropout=d.get("dropout", 0.0),
            lambda_weight=d.get("lambda", 1.0),
            flow=FlowConfig.from_dict(d.get("flow")),
            ddpm=DdpmConfig.from_dict(d.get("ddpm")),
            wgan=Stage1WganConfig.from_dict(d.get("wgan")),
            router=RouterConfig.from_dict(d.get("router")),
            trunk=TrunkConfig.from_dict(d.get("trunk")),
            heads=Stage1HeadsConfig.from_dict(d.get("heads")),
        )

    def to_dict(self) -> dict:
        return {
            "active": self.active,
            "init_from": self.init_from,
            "freeze": self.freeze,
            "generator": self.generator,
            "hidden_dim": self.hidden_dim,
            "n_res_blocks": self.n_res_blocks,
            "dropout": self.dropout,
            "lambda": self.lambda_weight,
            "flow": self.flow.to_dict(),
            "ddpm": self.ddpm.to_dict(),
            "wgan": self.wgan.to_dict(),
            "router": self.router.to_dict(),
            "trunk": self.trunk.to_dict(),
            "heads": self.heads.to_dict(),
        }


@dataclass(frozen=True)
class Stage2ModelConfig:
    # false trains stage 1 alone. giant rollout must then refuse the
    # checkpoint; giant predict still works.
    active: bool = True
    # See Stage1ModelConfig.init_from/.freeze — same semantics, this stage.
    init_from: str = ""
    freeze: bool = False
    # "one_shot": predict all k_max slots simultaneously with padded slots
    # masked from the loss (v0.2 behaviour).
    # "autoregressive": emit one secondary at a time in descending-energy
    # order.
    decoder: str = "autoregressive"
    # As stage1_model.generator, but under "autoregressive" this is the
    # objective for each token.
    generator: str = "wgan"
    hidden_dim: int = 256
    n_res_blocks: int = 6
    dropout: float = 0.0
    lambda_weight: float = 1.0  # dict key "lambda"
    # Maximum secondary slots. Under "one_shot" this is the fixed output
    # width; under "autoregressive" it is a safety cap on the generation loop.
    k_max: int = 15
    # Width of the projected stage-1 outcome fed into stage 2's conditioning.
    context_dim: int = 64
    # "truth": the ground-truth stage-1 target vector, detached — stage-level
    # teacher forcing (v0.2 behaviour). "sampled": stage 1's own sampled
    # output, closing the train/inference gap at the cost of a sampling pass
    # per batch and a moving target early in training.
    stage1_context: str = "truth"
    # Ramp for "sampled": P(condition on the ground-truth stage-1 outcome
    # rather than a fresh sample), linearly interpolated from ctx_p_start
    # (epoch 0) to ctx_p_end (the final epoch) — the same scheduled-sampling
    # shape as autoregressive.tf_p_start/tf_p_end, so stage 2 doesn't chase a
    # wildly moving stage-1 target in early epochs. Unread under "truth".
    ctx_p_start: float = 1.0
    ctx_p_end: float = 0.0
    n_sec: NSecConfig = field(default_factory=NSecConfig)
    particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
    autoregressive: AutoregressiveConfig = field(default_factory=AutoregressiveConfig)
    flow: FlowConfig = field(default_factory=FlowConfig)
    ddpm: DdpmConfig = field(default_factory=DdpmConfig)
    wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig)
    router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig)
    trunk: TrunkConfig = field(default_factory=TrunkConfig)
    heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
        d = d or {}
        return cls(
            active=d.get("active", True),
            init_from=d.get("init_from", ""),
            freeze=d.get("freeze", False),
            decoder=d.get("decoder", "autoregressive"),
            generator=d.get("generator", "wgan"),
            hidden_dim=d.get("hidden_dim", 256),
            n_res_blocks=d.get("n_res_blocks", 6),
            dropout=d.get("dropout", 0.0),
            lambda_weight=d.get("lambda", 1.0),
            k_max=d.get("k_max", 15),
            context_dim=d.get("context_dim", 64),
            stage1_context=d.get("stage1_context", "truth"),
            ctx_p_start=d.get("ctx_p_start", 1.0),
            ctx_p_end=d.get("ctx_p_end", 0.0),
            n_sec=NSecConfig.from_dict(d.get("n_sec")),
            particle_type=ParticleTypeConfig.from_dict(d.get("particle_type")),
            autoregressive=AutoregressiveConfig.from_dict(d.get("autoregressive")),
            flow=FlowConfig.from_dict(d.get("flow")),
            ddpm=DdpmConfig.from_dict(d.get("ddpm")),
            wgan=Stage2WganConfig.from_dict(d.get("wgan")),
            router=Stage2RouterConfig.from_dict(d.get("router")),
            trunk=TrunkConfig.from_dict(d.get("trunk")),
            heads=Stage2HeadsConfig.from_dict(d.get("heads")),
        )

    def to_dict(self) -> dict:
        return {
            "active": self.active,
            "init_from": self.init_from,
            "freeze": self.freeze,
            "decoder": self.decoder,
            "generator": self.generator,
            "hidden_dim": self.hidden_dim,
            "n_res_blocks": self.n_res_blocks,
            "dropout": self.dropout,
            "lambda": self.lambda_weight,
            "k_max": self.k_max,
            "context_dim": self.context_dim,
            "stage1_context": self.stage1_context,
            "ctx_p_start": self.ctx_p_start,
            "ctx_p_end": self.ctx_p_end,
            "n_sec": self.n_sec.to_dict(),
            "particle_type": self.particle_type.to_dict(),
            "autoregressive": self.autoregressive.to_dict(),
            "flow": self.flow.to_dict(),
            "ddpm": self.ddpm.to_dict(),
            "wgan": self.wgan.to_dict(),
            "router": self.router.to_dict(),
            "trunk": self.trunk.to_dict(),
            "heads": self.heads.to_dict(),
        }


@dataclass(frozen=True)
class TrainConfig:
    epochs: int = 100
    batch_size: int = 4096
    lr: float = 3e-4
    weight_decay: float = 0.01  # AdamW default — exposed so it can be tuned
    ema_decay: float = 0.9999  # EMA of model weights for sampling; 0 disables
    warmup_epochs: int = 5
    val_fraction: float = 0.1
    # per-epoch val loss (not the marginal/KL validate_every pass) is capped
    # to this many batches; 0 = full val set every epoch
    max_val_batches: int = 200
    num_workers: int = 4
    seed: int = 0
    validate_every: int = 10
    validate_steps: int = 10
    # Weights & Biases per-epoch metric logging. Default true in v0.3.0 (was
    # opt-in false): the v0.3.0 work is a sequence of architecture
    # comparisons, and a run that wasn't logged isn't comparable. Set false
    # for throwaway/debug runs.
    wandb: bool = True
    wandb_project: str = "giant"
    # "" means "use the checkpoint out_dir name" — not None, since the TOML
    # writer has no null literal to round-trip.
    wandb_run_name: str = ""
    # Batch-granularity metrics (loss/grad_norm/lr) are logged every N
    # optimizer steps, not every batch — a single epoch can be tens of
    # thousands of steps. Per-epoch metrics (the metrics.csv row) always log
    # in full.
    wandb_log_every: int = 50
    # Training-step autocast dtype: "fp32" (default, no autocast) or "bf16".
    # No "fp16" — GradScaler and the double-backward in
    # giant.model.wgan.gradient_penalty don't mix well, and bf16 alone covers
    # every training GPU in the fleet (Ampere and newer). See
    # giant.training.amp.resolve_autocast (gitea #47).
    precision: str = "fp32"

    @classmethod
    def from_dict(cls, d: dict | None) -> "TrainConfig":
        d = d or {}
        return cls(
            epochs=d.get("epochs", 100),
            batch_size=d.get("batch_size", 4096),
            lr=d.get("lr", 3e-4),
            weight_decay=d.get("weight_decay", 0.01),
            ema_decay=d.get("ema_decay", 0.9999),
            warmup_epochs=d.get("warmup_epochs", 5),
            val_fraction=d.get("val_fraction", 0.1),
            max_val_batches=d.get("max_val_batches", 200),
            num_workers=d.get("num_workers", 4),
            seed=d.get("seed", 0),
            validate_every=d.get("validate_every", 10),
            validate_steps=d.get("validate_steps", 10),
            wandb=d.get("wandb", True),
            wandb_project=d.get("wandb_project", "giant"),
            wandb_run_name=d.get("wandb_run_name", ""),
            wandb_log_every=d.get("wandb_log_every", 50),
            precision=d.get("precision", "fp32"),
        )

    def to_dict(self) -> dict:
        return {
            "epochs": self.epochs,
            "batch_size": self.batch_size,
            "lr": self.lr,
            "weight_decay": self.weight_decay,
            "ema_decay": self.ema_decay,
            "warmup_epochs": self.warmup_epochs,
            "val_fraction": self.val_fraction,
            "max_val_batches": self.max_val_batches,
            "num_workers": self.num_workers,
            "seed": self.seed,
            "validate_every": self.validate_every,
            "validate_steps": self.validate_steps,
            "wandb": self.wandb,
            "wandb_project": self.wandb_project,
            "wandb_run_name": self.wandb_run_name,
            "wandb_log_every": self.wandb_log_every,
            "precision": self.precision,
        }


@dataclass(frozen=True)
class GiantConfig:
    """Root config dataclass — the single source of truth for every default
    in DEFAULT_CONFIG below, which is generated from `GiantConfig().to_dict()`
    rather than hand-maintained (see issues.md Issue 1)."""

    conditioning: ConditioningConfig = field(default_factory=ConditioningConfig)
    stage1_model: Stage1ModelConfig = field(default_factory=Stage1ModelConfig)
    stage2_model: Stage2ModelConfig = field(default_factory=Stage2ModelConfig)
    train: TrainConfig = field(default_factory=TrainConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "GiantConfig":
        d = d or {}
        return cls(
            conditioning=ConditioningConfig.from_dict(d.get("conditioning")),
            stage1_model=Stage1ModelConfig.from_dict(d.get("stage1_model")),
            stage2_model=Stage2ModelConfig.from_dict(d.get("stage2_model")),
            train=TrainConfig.from_dict(d.get("train")),
        )

    def to_dict(self) -> dict:
        return {
            "conditioning": self.conditioning.to_dict(),
            "stage1_model": self.stage1_model.to_dict(),
            "stage2_model": self.stage2_model.to_dict(),
            "train": self.train.to_dict(),
        }


DEFAULT_CONFIG: dict = GiantConfig().to_dict()


def leaf_paths(node: dict, prefix: str = "") -> list[str]:
    """Every dotted leaf path in a DEFAULT_CONFIG-shaped dict, e.g.
    "stage1_model.router.n_experts". `[meta]` (run provenance, no schema
    counterpart) is skipped at the top level, matching `validate_config_keys`.
    Shared by `tests/test_config_consumed_keys.py` (the static per-identifier
    audit) and `giant.model.summary` (the runtime per-config audit, gitea
    #46) so both walk the exact same tree."""
    paths = []
    for key, value in node.items():
        if prefix == "" and key == "meta":
            continue
        path = f"{prefix}.{key}" if prefix else key
        if isinstance(value, dict):
            paths.extend(leaf_paths(value, path))
        else:
            paths.append(path)
    return paths


def git_hash() -> str:
    try:
        return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
    except Exception:
        return "unknown"


def auto_device() -> torch.device:
    import torch

    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


# Calibration point for estimate_batch_size(training=True): hidden_dim=1024,
# n_blocks=6, batch_size=29696 measured at ~7683 MiB VRAM (post-Phase-2
# architecture, including the Stage-2 secondary decoder and n_sec head).
# Activation memory is assumed to scale linearly with
# batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this
# is a rough estimate rather than a guaranteed bound.
# NOTE: not yet recalibrated for the v0.3.0 autoregressive stage-2 trunk —
# deliberately last in the implementation order.
_REF_BYTES = 7683 * 1024**2
_REF_BATCH_SIZE = 29696
_REF_HIDDEN_DIM = 1024
_REF_N_BLOCKS = 6

# Calibration point for estimate_batch_size(training=False): inference has no
# backward graph or optimizer state, so its memory footprint is much smaller
# per sample. hidden_dim=1024, n_blocks=8, batch_size=65536 measured at ~2037
# MiB VRAM.
_REF_BYTES_PREDICT = 2037 * 1024**2
_REF_BATCH_SIZE_PREDICT = 65536
_REF_HIDDEN_DIM_PREDICT = 1024
_REF_N_BLOCKS_PREDICT = 8


def estimate_batch_size(
    hidden_dim: int,
    n_blocks: int,
    device: torch.device,
    safety_factor: float = 0.8,
    min_batch_size: int = 1024,
    training: bool = True,
) -> int:
    """Estimate a batch size that fits in the free memory on `device`.

    Only supported on CUDA devices, which expose a free/total memory query;
    other backends (cpu, mps) raise ValueError. Pass `training=False` for
    inference (e.g. `predict`), which uses a much lower per-sample memory
    calibration since there's no backward graph or optimizer state.
    """
    import torch

    if device.type != "cuda":
        raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
    device_index = device.index if device.index is not None else torch.cuda.current_device()
    free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
    if training:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES,
            _REF_BATCH_SIZE,
            _REF_HIDDEN_DIM,
            _REF_N_BLOCKS,
        )
    else:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES_PREDICT,
            _REF_BATCH_SIZE_PREDICT,
            _REF_HIDDEN_DIM_PREDICT,
            _REF_N_BLOCKS_PREDICT,
        )
    bytes_per_unit = ref_bytes / (ref_batch_size * ref_hidden_dim * ref_n_blocks)
    bytes_per_sample = bytes_per_unit * hidden_dim * n_blocks
    batch_size = int(free_bytes * safety_factor / bytes_per_sample)
    batch_size = max(min_batch_size, (batch_size // 1024) * 1024)
    return batch_size


def load_toml(path: Path) -> dict:
    with open(path, "rb") as f:
        return tomllib.load(f)


def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None:
    """Warn (don't fail) if a config.toml's [meta].git_hash predates the current checkout.

    A config saved by a previous run may have been produced by code that has
    since changed, so its hyperparameters might not mean what they used to —
    surface that as a heads-up rather than blocking the rerun.
    """
    file_hash = file_cfg.get("meta", {}).get("git_hash")
    current_hash = git_hash()
    if not file_hash or file_hash == "unknown" or current_hash == "unknown":
        return
    if file_hash != current_hash:
        print(
            f"warning: {config_path} was generated at git commit {file_hash}, "
            f"but the current checkout is at {current_hash} — hyperparameters "
            "may not match the code that originally produced this config",
            file=sys.stderr,
        )


def load_checkpoint_config(ckpt_path: str | Path) -> dict:
    """Load the full config.toml written alongside a checkpoint by
    ``save_config``.

    Returns ``{}`` if no config.toml sits next to the checkpoint (older runs,
    or a checkpoint moved without its sidecar) — this is best-effort
    provenance for threading into a rollout's YAML sidecar, not a hard
    requirement for using the checkpoint itself. Returned as-loaded (v0.2 or
    v0.3 shape); callers that need the v0.3 shape should run it through
    `migrate_config` themselves.
    """
    config_path = Path(ckpt_path).parent / "config.toml"
    if not config_path.exists():
        return {}
    return load_toml(config_path)


def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None:
    """Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.

    Training writes config.toml into the same out_dir as its checkpoints, so a
    checkpoint loaded later (for `predict` or `giant.analysis`) can be
    cross-checked the same way `--config` loading is, without the caller having
    to pass the toml path explicitly. Silently does nothing if no config.toml
    is found alongside the checkpoint.
    """
    config_path = Path(ckpt_path).parent / "config.toml"
    if not config_path.exists():
        return
    warn_if_git_hash_mismatch(load_toml(config_path), config_path)


def _get_path(d: dict, dotted: str):
    """Read a dotted path (e.g. "stage1_model.router.enabled") out of a
    nested dict. Returns None if any component along the path is missing."""
    cur = d
    for part in dotted.split("."):
        if not isinstance(cur, dict) or part not in cur:
            return None
        cur = cur[part]
    return cur


def _set_path(d: dict, dotted: str, value) -> None:
    """Write a dotted path into a nested dict, creating intermediate dicts as
    needed."""
    parts = dotted.split(".")
    cur = d
    for part in parts[:-1]:
        cur = cur.setdefault(part, {})
    cur[parts[-1]] = value


def _pop_path(d: dict, dotted: str) -> None:
    """Remove a dotted path from a nested dict, if present. No-op if any
    component along the path is missing."""
    parts = dotted.split(".")
    cur = d
    for part in parts[:-1]:
        if not isinstance(cur, dict) or part not in cur:
            return
        cur = cur[part]
    if isinstance(cur, dict):
        cur.pop(parts[-1], None)


# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
# above) — normalized by migrate_config so a config.toml still using an older
# v0.3 key name keeps passing validate_config_keys.
_RENAMED_KEYS = {
    "stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling",  # gitea #86
}


def _deep_merge(base: dict, override: dict) -> dict:
    """Recursively merge `override` onto a copy of `base`.

    Dict-valued keys recurse instead of being replaced wholesale, so
    overriding one leaf (e.g. stage1_model.router.enabled) never drops
    untouched siblings — the rest of stage1_model.router, or of
    stage1_model — the same property v0.2's router-only bespoke merge had,
    generalized here to arbitrary depth.
    """
    result = dict(base)
    for k, v in override.items():
        if isinstance(v, dict) and isinstance(result.get(k), dict):
            result[k] = _deep_merge(result[k], v)
        else:
            result[k] = v
    return result


@dataclass(frozen=True)
class InferenceOverride:
    """One dotted `model_config` path that `giant.checkpoint_io.load_for_inference`
    is allowed to change on an already-trained checkpoint, without retraining.

    A path only belongs here if it affects neither module construction/tensor
    shapes nor the data preprocessing the normalizers/vocab maps were fit
    under — see the module docstring on `giant.model.summary` for the class
    of key this targets (`_fingerprint`'s "plain scalar attribute" leaves),
    and `giant.checkpoint_io.apply_config_overrides` for where this is used.
    """

    why: str
    choices: tuple[str, ...] | None = None
    minimum: float | None = None
    numeric: bool = False  # int/float leaf (vs. str, the default)

    def check(self, path: str, value: object) -> None:
        if self.choices is not None:
            if value not in self.choices:
                raise ValueError(f"{path} = {value!r} — must be one of {self.choices}")
            return
        if self.numeric:
            if isinstance(value, bool) or not isinstance(value, (int, float)):
                raise ValueError(f"{path} = {value!r} — must be a number")
            if self.minimum is not None and value < self.minimum:
                raise ValueError(f"{path} = {value!r} — must be >= {self.minimum}")


# Inference-safe dotted `model_config` paths — the allowlist gitea #87 asked
# for, so a typo or a shape-bearing key (e.g. "stage1_model.hidden_dim")
# raises a clear CheckpointCompatibilityError instead of surfacing as an
# opaque load_state_dict size mismatch later. Extend this table, not a
# per-call bypass, when a new inference-only key needs the capability.
INFERENCE_OVERRIDES: dict[str, InferenceOverride] = {
    "stage2_model.n_sec.sampling": InferenceOverride(
        why="giant.sample's n_sec head/stop-token sampling reads this at sample time only (gitea #86)",
        choices=STOP_SAMPLING_CHOICES,
    ),
    "stage1_model.ddpm.n_steps": InferenceOverride(
        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
        numeric=True,
        minimum=1,
    ),
    "stage2_model.ddpm.n_steps": InferenceOverride(
        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
        numeric=True,
        minimum=1,
    ),
    "stage2_model.particle_type.other_policy": InferenceOverride(
        why="giant.rollout resolves an 'other'-bucket secondary's PDG code with this at rollout time",
        choices=OTHER_POLICY_CHOICES,
    ),
    "stage1_model.router.temperature": InferenceOverride(
        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
        numeric=True,
        minimum=1e-6,
    ),
    "stage2_model.router.temperature": InferenceOverride(
        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
        numeric=True,
        minimum=1e-6,
    ),
}


@dataclass(frozen=True)
class FlagSpec:
    """One CLI flag's mapping into the config-overrides tree.

    `paths` lists every dotted config path this flag writes (>1 means fan-out
    to multiple stages/axes, e.g. `--mode` -> both stages' `generator`).
    `precedence` controls write order when two flags target the same path:
    specs are applied in ascending precedence, so a higher-precedence (more
    specific) flag overwrites a lower-precedence (shared/shorthand) one —
    this is the "build a shared dict, then let a more specific dict win"
    pattern `giant train`/`giant new-run` need (e.g. `--hidden-dim` vs
    `--stage1-hidden-dim`, or `--n-critic` vs `--stage1-n-critic`),
    generalized to one mechanism instead of three different ad hoc ones.
    """

    name: str
    paths: tuple[str, ...]
    precedence: int = 0


# Flag -> config-path table shared by `giant train`/`giant new-run`
# (giant/cli.py) so both commands resolve CLI overrides identically. See
# issues.md Issue 3: this replaces ~140 lines of hand-written, imperative
# dict-building in cli.py with one declarative table plus
# `overrides_from_flags` below.
FLAG_SPECS: tuple[FlagSpec, ...] = (
    # train block -- flat pass-through, unique paths, precedence irrelevant.
    FlagSpec("epochs", ("train.epochs",)),
    FlagSpec("batch_size", ("train.batch_size",)),
    FlagSpec("lr", ("train.lr",)),
    FlagSpec("weight_decay", ("train.weight_decay",)),
    FlagSpec("ema_decay", ("train.ema_decay",)),
    FlagSpec("warmup_epochs", ("train.warmup_epochs",)),
    FlagSpec("val_fraction", ("train.val_fraction",)),
    FlagSpec("num_workers", ("train.num_workers",)),
    FlagSpec("seed", ("train.seed",)),
    FlagSpec("validate_every", ("train.validate_every",)),
    FlagSpec("validate_steps", ("train.validate_steps",)),
    FlagSpec("max_val_batches", ("train.max_val_batches",)),
    FlagSpec("wandb", ("train.wandb",)),
    FlagSpec("wandb_project", ("train.wandb_project",)),
    FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
    FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
    FlagSpec("precision", ("train.precision",)),
    # --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat
    # shorthands (they predate stage2_model having its own flags);
    # --stage1-* wins when both are given.
    FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0),
    FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1),
    FlagSpec("n_blocks", ("stage1_model.n_res_blocks",), precedence=0),
    FlagSpec("stage1_n_res_blocks", ("stage1_model.n_res_blocks",), precedence=1),
    FlagSpec("dropout", ("stage1_model.dropout",), precedence=0),
    FlagSpec("stage1_dropout", ("stage1_model.dropout",), precedence=1),
    # stage2-only knobs.
    FlagSpec("stage2_hidden_dim", ("stage2_model.hidden_dim",)),
    FlagSpec("stage2_n_res_blocks", ("stage2_model.n_res_blocks",)),
    FlagSpec("stage2_dropout", ("stage2_model.dropout",)),
    FlagSpec("stage2_decoder", ("stage2_model.decoder",)),
    FlagSpec("stage2_k_max", ("stage2_model.k_max",)),
    FlagSpec("stage2_context_dim", ("stage2_model.context_dim",)),
    FlagSpec("stage2_stage1_context", ("stage2_model.stage1_context",)),
    # --mode applies to both stages by default (v0.2 had one shared
    # mode/wgan config); --stage{1,2}-generator override a single stage.
    FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator"), precedence=0),
    FlagSpec("stage1_generator", ("stage1_model.generator",), precedence=1),
    FlagSpec("stage2_generator", ("stage2_model.generator",), precedence=1),
    # --emb-dim/--conditioning set both conditioning axes (v0.2 had one
    # shared value for particle+material).
    FlagSpec("conditioning", ("conditioning.particle.type", "conditioning.material.type")),
    FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")),
    # Pre-aggregated router override dict (built by `_router_cli_overrides`
    # in cli.py from --router/--router-type/--n-experts/--router-axis).
    # Router overrides only ever land on stage1_model -- this asymmetry is
    # deliberate (see cli.py) and must not be "fixed" into a fan-out here.
    FlagSpec("router_config", ("stage1_model.router",)),
    # WGAN: shared knobs apply to both stages by default (v0.2 had one
    # shared wgan config); --stage{1,2}-* override a single stage.
    FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic"), precedence=0),
    FlagSpec("stage1_n_critic", ("stage1_model.wgan.n_critic",), precedence=1),
    FlagSpec("stage2_n_critic", ("stage2_model.wgan.n_critic",), precedence=1),
    FlagSpec("gp_weight", ("stage1_model.wgan.gp_weight", "stage2_model.wgan.gp_weight"), precedence=0),
    FlagSpec("stage1_gp_weight", ("stage1_model.wgan.gp_weight",), precedence=1),
    FlagSpec("stage2_gp_weight", ("stage2_model.wgan.gp_weight",), precedence=1),
    FlagSpec("noise_dim", ("stage1_model.wgan.noise_dim", "stage2_model.wgan.noise_dim"), precedence=0),
    FlagSpec("stage1_noise_dim", ("stage1_model.wgan.noise_dim",), precedence=1),
    FlagSpec("stage2_noise_dim", ("stage2_model.wgan.noise_dim",), precedence=1),
    FlagSpec("critic_lr", ("stage1_model.wgan.critic_lr", "stage2_model.wgan.critic_lr"), precedence=0),
    FlagSpec("stage1_critic_lr", ("stage1_model.wgan.critic_lr",), precedence=1),
    FlagSpec("stage2_critic_lr", ("stage2_model.wgan.critic_lr",), precedence=1),
    # Critic sizing: stage-scoped only, no shared alias — this is an
    # architectural per-stage knob like hidden_dim/n_res_blocks above, not a
    # shared training hyperparameter like the wgan knobs above it.
    FlagSpec("stage1_critic_hidden_dim", ("stage1_model.wgan.critic_hidden_dim",)),
    FlagSpec("stage1_critic_n_res_blocks", ("stage1_model.wgan.critic_n_res_blocks",)),
    FlagSpec("stage2_critic_hidden_dim", ("stage2_model.wgan.critic_hidden_dim",)),
    FlagSpec("stage2_critic_n_res_blocks", ("stage2_model.wgan.critic_n_res_blocks",)),
    # Partial-retrain (gitea #42): stage-scoped only, no shared alias — a
    # shared "freeze both stages from the same file" flag has no sensible
    # meaning (a checkpoint has one set of weights per stage).
    FlagSpec("stage1_init_from", ("stage1_model.init_from",)),
    FlagSpec("stage1_freeze", ("stage1_model.freeze",)),
    FlagSpec("stage2_init_from", ("stage2_model.init_from",)),
    FlagSpec("stage2_freeze", ("stage2_model.freeze",)),
)


def overrides_from_flags(values: dict[str, object]) -> dict:
    """Build the nested, section-keyed config-overrides dict
    `merge_cli_overrides` expects, from `{flag_name: value}`.

    Flags absent from `values`, or mapped to `None` (= not given on the
    CLI), are skipped. See `FlagSpec`/`FLAG_SPECS` above for the precedence
    rule applied when two flags target the same path.
    """
    overrides: dict = {}
    for spec in sorted(FLAG_SPECS, key=lambda s: s.precedence):
        if spec.name not in values or values[spec.name] is None:
            continue
        for path in spec.paths:
            _set_path(overrides, path, values[spec.name])
    return overrides


# v0.2 [train] keys that pass through to v0.3 [train] unchanged (same name,
# same meaning) when present in the loaded file — everything model-shaped
# moved to the stage/conditioning blocks instead (see the rest of
# migrate_config below).
_V02_TRAIN_PASSTHROUGH = (
    "epochs",
    "batch_size",
    "lr",
    "weight_decay",
    "ema_decay",
    "max_val_batches",
    "val_fraction",
    "num_workers",
    "seed",
    "validate_every",
    "validate_steps",
    "warmup_epochs",
    "wandb",
    "wandb_project",
    "wandb_run_name",
    "wandb_log_every",
)

# v0.2 train.{n_critic,gp_weight,critic_lr} applied identically to both
# stages' wgan sub-table (there was only ever one wgan objective, shared).
_V02_TRAIN_TO_BOTH_STAGES_WGAN = (
    ("n_critic", "n_critic"),
    ("gp_weight", "gp_weight"),
    ("critic_lr", "critic_lr"),
)


def migrate_config(cfg: dict) -> dict:
    """Translate a v0.2 config dict (single [train] + [model]) into the v0.3
    nested format ([conditioning]/[stage1_model]/[stage2_model]/[train]).

    Called on every config.toml load (see merge_cli_overrides) so old
    training configs on disk keep working under new code without hand-
    editing. `[meta].config_version == CONFIG_VERSION` marks a dict as
    already-v0.3; its absence is read as "this is v0.2", so an
    already-migrated dict is returned unchanged (deep-copied).

    Only keys actually present in `cfg` are translated — `cfg` may be a
    partial file (e.g. `[train]\\nepochs = 5\\n` with no [model] section at
    all, relying on v0.2 defaults for everything else). Separately, a fixed
    set of v0.2 architectural facts that were never exposed as config keys at
    all (e.g. the conditioning MLP was always 2 layers deep, not the v0.3
    default of 1) are injected unconditionally whenever this function decides
    it is migrating a v0.2 dict, regardless of which keys the file happened
    to set.

    Operates on the config.toml shape. A checkpoint's `model_config` dict
    (which additionally carries n_sec_head ownership and needs
    `network.build_models`'s cooperation) is a separate migration surface,
    deferred to the network.py refactor.

    Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
    keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
    gitea #86) so a config.toml written against an older v0.3 key name still
    passes `validate_config_keys`.
    """
    cfg = copy.deepcopy(cfg)
    for old_path, new_path in _RENAMED_KEYS.items():
        if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
            _set_path(cfg, new_path, _get_path(cfg, old_path))
        _pop_path(cfg, old_path)

    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
        return cfg

    old_train = cfg.pop("train", {})
    old_model = cfg.pop("model", {})
    old_router = dict(old_model.pop("router", {}))

    new: dict = {}

    for key in _V02_TRAIN_PASSTHROUGH:
        if key in old_train:
            _set_path(new, f"train.{key}", old_train[key])

    if "mode" in old_train:
        _set_path(new, "stage1_model.generator", old_train["mode"])
        _set_path(new, "stage2_model.generator", old_train["mode"])
    if "lambda_nsec" in old_train:
        _set_path(new, "stage2_model.n_sec.lambda", old_train["lambda_nsec"])
    if "lambda_s2" in old_train:
        _set_path(new, "stage2_model.lambda", old_train["lambda_s2"])
    for old_key, new_key in _V02_TRAIN_TO_BOTH_STAGES_WGAN:
        if old_key in old_train:
            _set_path(new, f"stage1_model.wgan.{new_key}", old_train[old_key])
            _set_path(new, f"stage2_model.wgan.{new_key}", old_train[old_key])

    for old_key, new_key in V02_MODEL_KEY_TO_STAGES:
        if old_key in old_model:
            _set_path(new, f"stage1_model.{new_key}", old_model[old_key])
            _set_path(new, f"stage2_model.{new_key}", old_model[old_key])
    if "emb_dim" in old_model:
        _set_path(new, "conditioning.particle.emb_dim", old_model["emb_dim"])
        _set_path(new, "conditioning.material.emb_dim", old_model["emb_dim"])
    if "conditioning" in old_model:
        _set_path(new, "conditioning.particle.type", old_model["conditioning"])
        _set_path(new, "conditioning.material.type", old_model["conditioning"])
    if "noise_dim" in old_model:
        _set_path(new, "stage1_model.wgan.noise_dim", old_model["noise_dim"])
        _set_path(new, "stage2_model.wgan.noise_dim", old_model["noise_dim"])
    if "k_max" in old_model:
        _set_path(new, "stage2_model.k_max", old_model["k_max"])

    if old_router:
        reject_legacy_router_expert_sizing(old_router, source="v0.2 config's model.router")
        _set_path(new, "stage1_model.router", dict(old_router))
        stage2_router = dict(old_router)
        stage2_router["tie_to_stage1"] = False
        _set_path(new, "stage2_model.router", stage2_router)

    # v0.2 architectural facts with no corresponding config key at all —
    # always set once we've determined we're migrating a v0.2 dict,
    # independent of what the file did/didn't specify (see giant._migration).
    for path, value in V02_FIXED_FACTS.items():
        _set_path(new, path, value)

    new_meta = dict(cfg.pop("meta", {}))
    new_meta["config_version"] = CONFIG_VERSION
    new["meta"] = new_meta

    # Anything else in the original dict (unrecognized top-level sections)
    # carries through untouched rather than being silently dropped.
    for k, v in cfg.items():
        new.setdefault(k, v)

    return new


# axis{i}_{field} composed-router keys (see network._parse_composed_axes) — field-name
# agnostic, matching network.py's own _AXIS_KEY_RE, since anything after axis{i}_ is
# passed straight through as a router kwarg there.
_AXIS_KEY_RE = re.compile(r"^axis\d+_.+$")

# Paths (dotted, relative to the merged cfg root) that carry genuinely dynamic keys not
# in DEFAULT_CONFIG's fixed schema — composed-router axis{i}_{field} flags and
# pipeline.seed_router_centers's runtime-seeded centers_init (see RouterConfig.extra
# above). validate_config_keys allows any key under these paths through unconditionally
# apart from the axis-pattern/centers_init check below.
_DYNAMIC_ROUTER_PATHS = {"stage1_model.router", "stage2_model.router"}


def _unknown_key_error(full_path: str, key: str, valid_keys) -> ValueError:
    hint = difflib.get_close_matches(key, list(valid_keys), n=1)
    suggestion = f" — did you mean {hint[0]!r}?" if hint else ""
    return ValueError(f"unknown config key {full_path!r}{suggestion}")


def _validate_keys(node: dict, default_node: dict, path: str) -> None:
    for key, value in node.items():
        if path == "" and key == "meta":
            continue
        full_path = f"{path}.{key}" if path else key
        if path in _DYNAMIC_ROUTER_PATHS and (key == "centers_init" or _AXIS_KEY_RE.match(key)):
            continue
        if key not in default_node:
            raise _unknown_key_error(full_path, key, default_node.keys())
        if isinstance(value, dict) and isinstance(default_node[key], dict):
            _validate_keys(value, default_node[key], full_path)


def validate_config_keys(cfg: dict) -> None:
    """Reject any config key not part of the known v0.3 schema (DEFAULT_CONFIG's tree).

    Catches typos like `n_res_block` for `n_res_blocks` that would otherwise merge
    cleanly, pass `validate_config`, and silently build the wrong model — see
    issues.md Issue 2.

    Only exercised on the config.toml/CLI-overrides path (called from
    `merge_cli_overrides` below). A checkpoint's `model_config` dict goes through
    `network._migrate_legacy_model_config`/`build_models` instead and must keep
    loading regardless of schema drift; old checkpoints predate this validator and
    are never passed through here.

    Two areas are deliberately dynamic and excluded: `[meta]` (run provenance, no
    DEFAULT_CONFIG counterpart), and `stage{1,2}_model.router`'s `axis{i}_{field}`
    keys (composed-router axes, see `network._parse_composed_axes`) /
    `centers_init` (runtime-seeded by `pipeline.seed_router_centers`).
    """
    _validate_keys(cfg, DEFAULT_CONFIG, "")


def merge_cli_overrides(
    defaults: dict,
    config_path: Path | None,
    overrides: dict,
) -> dict:
    """Resolve config as defaults -> TOML file -> explicit overrides.

    `overrides` is keyed by top-level section name (e.g. "stage1_model",
    "train"), each value an arbitrarily nested dict of overrides to
    deep-merge (see `_deep_merge`) — the shape stage-prefixed CLI flags
    naturally produce. A v0.2-shaped TOML file is transparently migrated
    (`migrate_config`) before merging, so old configs on disk keep working
    under the new schema. The result is checked against the known schema
    (`validate_config_keys`) before being returned, so a typo'd key raises
    here rather than silently building the wrong model.
    """
    cfg = copy.deepcopy(defaults)
    if config_path is not None:
        file_cfg = migrate_config(load_toml(config_path))
        for section, values in file_cfg.items():
            if section == "meta":
                continue
            if isinstance(values, dict):
                cfg[section] = _deep_merge(cfg.get(section, {}), values)
            else:
                cfg[section] = values
        warn_if_git_hash_mismatch(file_cfg, config_path)
    for section, values in overrides.items():
        if isinstance(values, dict):
            cfg[section] = _deep_merge(cfg.get(section, {}), values)
        else:
            cfg[section] = values
    validate_config_keys(cfg)
    return cfg


def validate_config(cfg: dict, *, resume: bool = False) -> None:
    """Cross-block validation the per-block schema can't express on its own.

    Raises ValueError with a clear message on the first violation found. Call
    after `merge_cli_overrides` has produced a fully-merged v0.3 config —
    these checks need to see across blocks, so they don't belong in
    `migrate_config` (which only ever sees one dict's own keys) or in any
    single block's defaults.

    `resume=True` (only `giant train --resume` passes this) relaxes the
    `stage{1,2}_model.freeze` -> `.init_from` requirement below: a resumed
    frozen stage's weights come from the resume checkpoint, not `init_from`.
    """
    particle_type = _get_path(cfg, "conditioning.particle.type")

    pt_target = _get_path(cfg, "stage2_model.particle_type.target")
    if pt_target == "embedding" and particle_type != "embedding":
        raise ValueError(
            "stage2_model.particle_type.target = 'embedding' requires "
            "conditioning.particle.type = 'embedding' (there is no embedding "
            "table to regress against under conditioning.particle.type = "
            f"{particle_type!r})"
        )

    class_weighting = _get_path(cfg, "stage2_model.particle_type.class_weighting")
    if class_weighting not in ("none", "inverse_freq"):
        raise ValueError(
            f"stage2_model.particle_type.class_weighting = {class_weighting!r} — must be 'none' or 'inverse_freq'"
        )
    if class_weighting != "none" and pt_target != "onehot":
        raise ValueError(
            "stage2_model.particle_type.class_weighting != 'none' requires "
            f"stage2_model.particle_type.target = 'onehot' (there is no class "
            f"cross-entropy to weight under target = {pt_target!r})"
        )
    if class_weighting != "none" and _get_path(cfg, "stage2_model.generator") == "wgan":
        raise ValueError(
            "stage2_model.particle_type.class_weighting != 'none' is "
            "incompatible with stage2_model.generator = 'wgan' — that path "
            "feeds the type slice to the critic via a straight-through "
            "Gumbel relaxation instead of a class cross-entropy, so there is "
            "nothing to weight"
        )

    for stage_name in ("stage1_model", "stage2_model"):
        if _get_path(cfg, f"{stage_name}.freeze") and not _get_path(cfg, f"{stage_name}.init_from") and not resume:
            raise ValueError(
                f"{stage_name}.freeze = true requires {stage_name}.init_from "
                "to be set (or --resume) — freezing a randomly-initialized "
                "model is almost certainly a mistake"
            )

        router = _get_path(cfg, f"{stage_name}.router") or {}
        if router.get("enabled") and router.get("type") in ("pdg", "process") and particle_type == "physical":
            raise ValueError(
                f"{stage_name}.router.type = {router['type']!r} builds its "
                "own training-vocab-scoped embedding, incompatible with "
                "conditioning.particle.type = 'physical' (defeats "
                "generalization beyond the training menu) — pick a "
                "different router type or a different "
                "conditioning.particle.type"
            )

    if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path(cfg, "stage1_model.active"):
        raise ValueError(
            "stage2_model.router.tie_to_stage1 = true requires "
            "stage1_model.active = true (there is no stage-1 router to tie to)"
        )

    if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token":
        if _get_path(cfg, "stage2_model.decoder") != "autoregressive":
            raise ValueError(
                "stage2_model.n_sec.mode = 'stop_token' requires "
                "stage2_model.decoder = 'autoregressive' — there is no "
                "per-token loop to stop under 'one_shot'"
            )
        if _get_path(cfg, "stage2_model.n_sec.owner") != "stage2":
            raise ValueError(
                "stage2_model.n_sec.mode = 'stop_token' requires "
                "stage2_model.n_sec.owner = 'stage2' — a migrated v0.2 "
                "checkpoint's stage-1 n_sec_head has no per-token "
                "conditioning to hang an EOS decision off"
            )

    n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
    if n_sec_sampling not in STOP_SAMPLING_CHOICES:
        raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")

    precision = _get_path(cfg, "train.precision")
    if precision not in ("fp32", "bf16"):
        raise ValueError(
            f"train.precision = {precision!r} — must be 'fp32' or 'bf16' "
            "('fp16' is not supported: see giant.training.amp.resolve_autocast)"
        )

    stage1_context = _get_path(cfg, "stage2_model.stage1_context")
    if stage1_context not in ("truth", "sampled"):
        raise ValueError(f"stage2_model.stage1_context = {stage1_context!r} — must be 'truth' or 'sampled'")
    if stage1_context == "sampled":
        if not (_get_path(cfg, "stage1_model.active") and _get_path(cfg, "stage2_model.active")):
            raise ValueError(
                "stage2_model.stage1_context = 'sampled' requires both "
                "stage1_model.active and stage2_model.active = true — there is "
                "no stage-1 model to sample from in a stage-2-only run"
            )
        ctx_p_start = _get_path(cfg, "stage2_model.ctx_p_start")
        ctx_p_end = _get_path(cfg, "stage2_model.ctx_p_end")
        for name, value in (("ctx_p_start", ctx_p_start), ("ctx_p_end", ctx_p_end)):
            if not (0.0 <= value <= 1.0):
                raise ValueError(f"stage2_model.{name} = {value} — must be in [0, 1]")
        if ctx_p_start == 1.0 and ctx_p_end == 1.0:
            raise ValueError(
                "stage2_model.stage1_context = 'sampled' with ctx_p_start = "
                "ctx_p_end = 1.0 always conditions on the ground truth — "
                "identical to 'truth' but silently so; use 'truth' instead or "
                "lower ctx_p_end"
            )

    if (
        _get_path(cfg, "stage2_model.n_sec.mode") == "truth"
        and _get_path(cfg, "stage1_model.active")
        and _get_path(cfg, "stage2_model.active")
    ):
        raise ValueError(
            "stage2_model.n_sec.mode = 'truth' is invalid for a "
            "rollout-capable checkpoint (both stage1_model.active and "
            "stage2_model.active = true): "
            "'truth' takes n_sec from ground truth, which giant rollout "
            "doesn't have. 'truth' is for standalone stage-2 evaluation "
            "only — set stage1_model.active = false for that, or use "
            "'head' (default) for a rollout-capable checkpoint."
        )

    if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
        order = _get_path(cfg, "stage2_model.autoregressive.order")
        if order != "energy_desc":
            raise ValueError(
                f"stage2_model.autoregressive.order = {order!r} — must be "
                "'energy_desc' (the only implemented ordering; see "
                "AutoregressiveConfig.order's docstring)"
            )
        from giant.model.history import HISTORY_REGISTRY

        history = _get_path(cfg, "stage2_model.autoregressive.history")
        if history not in HISTORY_REGISTRY:
            raise ValueError(
                f"stage2_model.autoregressive.history = {history!r} — must be one of {sorted(HISTORY_REGISTRY)}"
            )
        teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
        if teacher_forcing not in ("always", "scheduled", "never"):
            raise ValueError(
                "stage2_model.autoregressive.teacher_forcing = "
                f"{teacher_forcing!r} — must be 'always', 'scheduled' or "
                "'never'"
            )


_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb", "onehot": "oh"}


def _path_candidate(dotted_path: str, prefix: str, formatter=str):
    """Candidate factory: show `prefix + formatter(value)` when the value at
    `dotted_path` differs from its DEFAULT_CONFIG value, else omit."""

    def _candidate(cfg):
        value = _get_path(cfg, dotted_path)
        default = _get_path(DEFAULT_CONFIG, dotted_path)
        if value == default:
            return None
        return f"{prefix}{formatter(value)}"

    return _candidate


def _conditioning_candidate(axis: str, short: str):
    def _candidate(cfg):
        value = _get_path(cfg, f"conditioning.{axis}.type")
        default = _get_path(DEFAULT_CONFIG, f"conditioning.{axis}.type")
        if value == default:
            return None
        code = _CONDITIONING_CODE.get(value, value)
        return f"{short}{code}"

    return _candidate


def _router_candidate(stage_key: str, short: str):
    """Candidate for a stage's router as a single unit, gated on
    `router.enabled` (a disabled router's type/n_experts are meaningless)."""

    def _candidate(cfg):
        router = _get_path(cfg, f"{stage_key}.router") or {}
        default_router = _get_path(DEFAULT_CONFIG, f"{stage_key}.router") or {}
        if router.get("enabled") == default_router.get("enabled"):
            return None
        return f"{short}r-{router['type']}{router['n_experts']}"

    return _candidate


def _router_flag_candidate(stage_key: str, short: str, field: str, token_map: dict):
    """Candidate factory for a boolean field inside a stage's router block.

    Gated on `router.enabled` like `_router_candidate`, then omitted unless
    `field` differs from its DEFAULT_CONFIG value. `token_map` need only
    cover the non-default value(s), since the default value always yields
    None.
    """

    def _candidate(cfg):
        router = _get_path(cfg, f"{stage_key}.router") or {}
        default_router = _get_path(DEFAULT_CONFIG, f"{stage_key}.router") or {}
        if router.get("enabled") == default_router.get("enabled"):
            return None
        value = router.get(field)
        if value == default_router.get(field):
            return None
        return f"{short}{token_map[value]}"

    return _candidate


# Priority-ordered candidate fields for default_out_dir_name: (label,
# candidate(cfg) -> str | None). Beyond _OUT_DIR_NAME_MAX_FIELDS non-default
# fields, the remainder collapse into a hash suffix (see
# default_out_dir_name). Candidates read the whole nested cfg via dotted
# paths — there is no single "model" dict anymore now that architecture is
# split across conditioning/stage1_model/stage2_model.
_OUT_DIR_NAME_CANDIDATES = [
    ("stage1_generator", _path_candidate("stage1_model.generator", "")),
    ("stage2_generator", _path_candidate("stage2_model.generator", "s2-")),
    ("stage2_decoder", _path_candidate("stage2_model.decoder", "dec-")),
    (
        "stage2_history",
        _path_candidate("stage2_model.autoregressive.history", "hist-"),
    ),
    (
        "particle_type_target",
        _path_candidate("stage2_model.particle_type.target", "pt-"),
    ),
    ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")),
    ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")),
    ("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")),
    ("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")),
    ("stage1_router", _router_candidate("stage1_model", "s1")),
    ("stage2_router", _router_candidate("stage2_model", "s2")),
    (
        "stage1_gumbel",
        _router_flag_candidate("stage1_model", "s1", "gumbel", {True: "gum"}),
    ),
    (
        "stage2_gumbel",
        _router_flag_candidate("stage2_model", "s2", "gumbel", {True: "gum"}),
    ),
    (
        "stage1_learn_centers",
        _router_flag_candidate("stage1_model", "s1", "learn_centers", {False: "nolc"}),
    ),
    (
        "stage2_learn_centers",
        _router_flag_candidate("stage2_model", "s2", "learn_centers", {False: "nolc"}),
    ),
    (
        "stage1_learn_width",
        _router_flag_candidate("stage1_model", "s1", "learn_width", {True: "lw"}),
    ),
    (
        "stage2_learn_width",
        _router_flag_candidate("stage2_model", "s2", "learn_width", {True: "lw"}),
    ),
    (
        "stage1_learn_temperature",
        _router_flag_candidate("stage1_model", "s1", "learn_temperature", {True: "lt"}),
    ),
    (
        "stage2_learn_temperature",
        _router_flag_candidate("stage2_model", "s2", "learn_temperature", {True: "lt"}),
    ),
    ("particle_conditioning", _conditioning_candidate("particle", "c")),
    ("material_conditioning", _conditioning_candidate("material", "m")),
    ("stage1_freeze", _path_candidate("stage1_model.freeze", "s1frozen", formatter=lambda _: "")),
    ("stage2_freeze", _path_candidate("stage2_model.freeze", "s2frozen", formatter=lambda _: "")),
    ("stage1_hidden_dim", _path_candidate("stage1_model.hidden_dim", "h")),
    ("stage2_hidden_dim", _path_candidate("stage2_model.hidden_dim", "s2h")),
    ("stage1_n_res_blocks", _path_candidate("stage1_model.n_res_blocks", "b")),
    ("stage2_n_res_blocks", _path_candidate("stage2_model.n_res_blocks", "s2b")),
    ("particle_emb_dim", _path_candidate("conditioning.particle.emb_dim", "e")),
    ("lr", _path_candidate("train.lr", "lr")),
    ("batch_size", _path_candidate("train.batch_size", "bs")),
    ("seed", _path_candidate("train.seed", "seed")),
    ("epochs", _path_candidate("train.epochs", "ep")),
]

_OUT_DIR_NAME_MAX_FIELDS = 6


def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
    """Build a default checkpoint out_dir name from what's non-default in `cfg`.

    Only fields that differ from DEFAULT_CONFIG are included, so a fully
    default run's name is just its timestamp — see
    `_OUT_DIR_NAME_CANDIDATES` for the fixed, priority-ordered field list.
    Beyond `_OUT_DIR_NAME_MAX_FIELDS` non-default fields, the remainder
    collapse into a short deterministic hash suffix rather than growing the
    name unboundedly. This name doubles as the run's W&B id (see
    giant.training), which is the reason a timestamp is always included.
    """
    now = now or datetime.now()
    tokens = []
    overflow = []
    for label, candidate in _OUT_DIR_NAME_CANDIDATES:
        token = candidate(cfg)
        if token is None:
            continue
        if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS:
            tokens.append(token)
        else:
            overflow.append(f"{label}={token}")

    name = now.strftime("%Y%m%d_%H%M")
    if tokens:
        name += "_" + "_".join(tokens)
    if overflow:
        digest = hashlib.md5("|".join(sorted(overflow)).encode()).hexdigest()[:6]
        name += f"_+{len(overflow)}more-{digest}"
    return name


def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path:
    """Auto-derived out dir from cfg's hyperparams (see `default_out_dir_name`),
    with a numeric suffix loop so two runs whose name collides (same
    non-default hyperparams, same to-the-minute timestamp) don't clobber each
    other's directory. Shared by `giant train` and `giant new-run`.
    """
    base_name = default_out_dir_name(cfg)
    out_dir = base / base_name
    suffix = 2
    while out_dir.exists():
        out_dir = base / f"{base_name}_{suffix}"
        suffix += 1
    return out_dir


def seed_everything(seed: int) -> None:
    import numpy as np
    import torch

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _toml_value(v) -> str:
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, str):
        return repr(v)
    return str(v)


def _write_section(lines: list[str], path: str, values: dict) -> None:
    """Write one TOML table (`[path]`) and recurse depth-first into any
    dict-valued keys as `[path.subkey]` — handles the v0.3 schema's 2-3 level
    nesting (e.g. stage1_model.router, stage2_model.n_sec) with no depth
    limit, unlike the one-level-only writer this replaces."""
    lines.append(f"[{path}]")
    nested: list[tuple[str, dict]] = []
    for k, v in values.items():
        if isinstance(v, dict):
            nested.append((f"{path}.{k}", v))
        else:
            lines.append(f"{k:<18} = {_toml_value(v)}")
    lines.append("")
    for sub_path, sub_values in nested:
        _write_section(lines, sub_path, sub_values)


def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
    lines: list[str] = []
    for section, values in cfg.items():
        _write_section(lines, section, values)

    lines.append("[meta]")
    for k, v in meta.items():
        lines.append(f"{k:<18} = {_toml_value(v)}")

    (out_dir / "config.toml").write_text("\n".join(lines))


def build_run_meta(
    data: Path,
    seed: int,
    n_pdg_codes: int,
    n_materials: int,
    n_train_events: int,
    n_val_events: int,
    n_train_steps: int,
) -> dict:
    import torch

    return {
        "config_version": CONFIG_VERSION,
        "git_hash": git_hash(),
        "seed": seed,
        "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "command": " ".join(sys.argv),
        "data_path": str(data),
        "n_pdg_codes": n_pdg_codes,
        "n_materials": n_materials,
        "n_train_events": n_train_events,
        "n_val_events": n_val_events,
        "n_train_steps": n_train_steps,
    }