giant/giant/model/network.py

import inspect
import math
import re

import torch
import torch.nn as nn
import torch.nn.functional as F

from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM


class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        assert dim % 2 == 0, "dim must be even"
        half = dim // 2
        freqs = torch.exp(
            -math.log(10000)
            * torch.arange(half, dtype=torch.float32)
            / max(half - 1, 1)
        )
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t = t.reshape(-1, 1).float()
        args = t * self.freqs.unsqueeze(0)  # (B, half)
        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)


class ConditionEncoder(nn.Module):
    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        cont_dim: int = COND_DIM,
        emb_dim: int = 16,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        in_dim = cont_dim + 2 * emb_dim
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, out_dim),
            nn.SiLU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        x = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
        return self.mlp(x)


class ResBlock(nn.Module):
    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear1 = nn.Linear(dim, dim)
        self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
        self.act = nn.SiLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        h = self.linear1(h) + self.cond_proj(cond)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + h


class DenoisingMLP(nn.Module):
    """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits.

    The n_sec head runs on the condition encoding only (no diffusion noise),
    so it can be called at inference time independently via `predict_n_sec`.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        k_max: int = K_MAX,
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(x_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        # Predicts n_sec as classification over {0, 1, ..., k_max}.
        # Applied to the condition encoding (not the diffused latent).
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, hidden_dim // 2),
            nn.SiLU(),
            nn.Linear(hidden_dim // 2, k_max + 1),
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)  # (B, time_dim)
        c_emb = self.cond_enc(cond_cont, cond_cat)  # (B, cond_out_dim)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)

    def predict_n_sec(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)

    def pdg_embedding_weight(self) -> torch.Tensor:
        """Return the PDG embedding table weights for secondary type targets."""
        return self.cond_enc.pdg_emb.weight


class SecondaryConditionEncoder(nn.Module):
    """Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_dim: int = X_DIM,
        stage1_proj_dim: int = 64,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.base = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
        fused_dim = cond_out_dim + stage1_proj_dim
        self.fuse = nn.Sequential(
            nn.Linear(fused_dim, out_dim),
            nn.SiLU(),
        )

    def forward(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        base = self.base(cond_cont, cond_cat)  # (B, cond_out_dim)
        s1 = self.stage1_proj(stage1_out).tanh()  # (B, stage1_proj_dim)
        return self.fuse(torch.cat([base, s1], dim=-1))  # (B, out_dim)


class SecondaryDecoder(nn.Module):
    """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously.

    Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one
    secondary ordered by descending energy. Padded slots are masked from loss.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


class Router(nn.Module):
    """Contract for a pluggable mixture-of-experts routing axis.

    Subclasses implement `gate` (soft partition-of-unity weights over
    experts, used in train mode for a fully differentiable mixture);
    `top1` and `balance_loss` have working defaults so a new routing axis
    is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
    """

    def __init__(self, n_experts: int) -> None:
        super().__init__()
        self.n_experts = n_experts

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B, n_experts) soft weights, rows summing to 1."""
        raise NotImplementedError

    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B,) hard expert index, used for eval-time grouped dispatch."""
        return self.gate(cond_cont, cond_cat).argmax(dim=-1)

    def balance_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
        importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
        return (importance.std() / (importance.mean() + 1e-8)) ** 2

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        """Optional supervised auxiliary loss shaping the router's own belief.

        Default: none (a scalar 0), for routers like EnergyRouter that read a
        quantity directly off cond_cont/cond_cat and need no label. Routers
        gating on an unobservable pre-step quantity (e.g. ProcessRouter,
        which predicts the physics process that will end the step) override
        this to supervise their internal classifier against the true label.
        """
        return torch.zeros((), device=cond_cont.device)


ROUTER_REGISTRY: dict[str, type[Router]] = {}


def register_router(name: str):
    def decorator(cls: type[Router]) -> type[Router]:
        ROUTER_REGISTRY[name] = cls
        return cls

    return decorator


def build_router(name: str, n_experts: int, **kwargs) -> Router:
    """Factory: look up a `Router` subclass by name from the registry.

    Every registered router type is fed the same `model.router` config
    dict; kwargs not declared by that type's constructor are silently
    dropped, so per-type hyperparameters (e.g. EnergyRouter's
    `temperature`) can coexist in one config without special-casing.
    """
    if name not in ROUTER_REGISTRY:
        raise ValueError(
            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
        )
    cls = ROUTER_REGISTRY[name]
    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
    filtered = {k: v for k, v in kwargs.items() if k in accepted}
    return cls(n_experts=n_experts, **filtered)


@register_router("energy")
class EnergyRouter(Router):
    """Soft turn-on gate over normalized pre-step log-energy.

    Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or
    fixed) 1-D centers, initialized spread across [-2, 2] — roughly the
    z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`,
    differentiable in e; as tau -> 0 this hardens to nearest-center
    (Voronoi) selection, which is exactly what `top1` uses at eval.
    """

    def __init__(
        self,
        n_experts: int = 4,
        temperature: float = 0.5,
        learn_centers: bool = True,
        energy_idx: int = 3,
    ) -> None:
        super().__init__(n_experts)
        self.temperature = temperature
        self.energy_idx = energy_idx
        centers = torch.linspace(-2.0, 2.0, n_experts)
        if learn_centers:
            self.centers = nn.Parameter(centers)
        else:
            self.register_buffer("centers", centers)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        e = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
        d2 = (e - self.centers.unsqueeze(0)) ** 2  # (B, n_experts)
        return torch.softmax(-d2 / self.temperature, dim=-1)


@register_router("pdg")
class PdgRouter(Router):
    """Soft turn-on gate over a learned PDG embedding.

    Unlike ProcessRouter's process label, PDG code is already known at
    pre-step time (it's a conditioning input, `cond_cat[:, 0]`), so no
    supervision is needed — `classify_loss` falls back to the Router base
    class's zero-loss default, same as EnergyRouter. Because PDG is
    categorical rather than a scalar, this generalizes EnergyRouter's
    soft-turn-on-then-Voronoi trick from a 1-D distance to a distance in a
    small embedding space: its own embedding table (kept separate from the
    trunk's ConditionEncoder, same reasoning as ProcessRouter's own
    pdg/mat embeddings) maps each PDG code to a point, and `n_experts`
    learnable (or fixed) centers partition that space.
    `gate(pdg) = softmax_i(-||emb(pdg) - c_i||^2 / tau)`.
    """

    def __init__(
        self,
        n_experts: int,
        pdg_vocab: int,
        emb_dim: int = 8,
        temperature: float = 0.5,
        learn_centers: bool = True,
    ) -> None:
        super().__init__(n_experts)
        self.temperature = temperature
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        centers = torch.randn(n_experts, emb_dim) * 0.1
        if learn_centers:
            self.centers = nn.Parameter(centers)
        else:
            self.register_buffer("centers", centers)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
            -1
        )  # (B, n_experts)
        return torch.softmax(-d2 / self.temperature, dim=-1)


@register_router("process")
class ProcessRouter(Router):
    """Routes on the physics process expected to end the step.

    Unlike EnergyRouter (which reads a quantity that's already known at
    pre-step time), the process — Compton, photoelectric, brems, ... — is a
    *post-step outcome*: it can't be read off cond_cont/cond_cat directly.
    Instead this router runs a small classifier over pre-step conditioning
    (its own pdg/material embeddings, kept separate from the trunk's
    ConditionEncoder) that predicts it, one class per expert slot
    (`n_experts` doubles as the number of process classes — see
    `build_process_map_from_files`, which caps the process vocabulary to
    exactly this many classes, bucketing rare processes into a shared
    "other" slot).

    The classifier is supervised by `classify_loss` against the true
    `process` label (see `giant/train.py`) — a *training-time* signal only;
    `gate`/`top1` never see it, so eval-time dispatch (rollout, predict)
    needs no ground truth, same as every other Router. This sidesteps the
    gradient/differentiability problem that sank the earlier
    process-conditioned-flow proposal (see the archived decision doc): the
    hard categorical choice only ever feeds a non-differentiable expert
    *dispatch*, never the flow's own conditioning path.
    """

    def __init__(
        self,
        n_experts: int,
        pdg_vocab: int,
        mat_vocab: int,
        emb_dim: int = 8,
        hidden_dim: int = 64,
    ) -> None:
        super().__init__(n_experts)
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        self.classifier = nn.Sequential(
            nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, n_experts),
        )

    def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B, n_experts) raw process-classifier logits, one class per expert."""
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
        return self.classifier(h)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)


class ComposedRouter(Router):
    """Joint router over independent axes (e.g. energy x pdg), outer-product gated.

    Wraps N already-built sub-routers, each free to have its own
    `n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)`
    composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match
    the other's expert count). The joint gate is the outer product of the
    per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a
    partition of unity, since each factor is one. Because the axes are
    routed independently, the joint argmax factors into the per-axis
    argmaxes, so `top1` (inherited from `Router`) costs no more than
    routing each axis alone despite the multiplicative expert count; the
    same is true of `balance_loss` (inherited, computed on the flattened
    joint gate — now one importance term per *joint* expert cell).

    Not registered in `ROUTER_REGISTRY` / buildable via `build_router`,
    since those assume one `n_experts` int shared by a single router type;
    use `build_composed_router` instead, which resolves a list of per-axis
    specs (each independently typed and sized) through `build_router`.
    """

    def __init__(self, routers: list[Router]) -> None:
        if not routers:
            raise ValueError("ComposedRouter needs at least one sub-router")
        n_experts = 1
        for r in routers:
            n_experts *= r.n_experts
        super().__init__(n_experts)
        self.routers = nn.ModuleList(routers)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
        for router in self.routers[1:]:
            g = router.gate(cond_cont, cond_cat)  # (B, n_i)
            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
                1
            )  # (B, prod so far)
        return joint

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        """Sum of each sub-router's own classify_loss (0 for unsupervised axes)."""
        total = torch.zeros((), device=cond_cont.device)
        for router in self.routers:
            total = total + router.classify_loss(cond_cont, cond_cat, labels)
        return total


def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
    """Build a `ComposedRouter` from a list of per-axis router specs.

    Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}`
    dict resolved through `build_router` exactly like a single-axis router
    config, so axes can differ in both expert count and hyperparameters
    (e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`).
    `shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each
    spec, with the spec's own keys taking precedence.
    """
    routers = [
        build_router(
            spec["type"],
            spec["n_experts"],
            **{
                **shared_kwargs,
                **{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
            },
        )
        for spec in specs
    ]
    return ComposedRouter(routers)


class ExpertTrunk(nn.Module):
    """One small expert: `input_proj -> ResBlock stack -> out_proj`.

    Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but
    intended to be narrower/shallower (per-call cost is the whole point).
    """

    def __init__(
        self,
        in_dim: int,
        hidden_dim: int,
        n_blocks: int,
        merged_cond_dim: int,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.input_proj = nn.Linear(in_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, in_dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


def _route_forward(
    experts: nn.ModuleList,
    router: Router,
    x: torch.Tensor,
    cond: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    training: bool,
) -> torch.Tensor:
    """Shared dispatch for both Routed* trunks.

    Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully
    differentiable, N-expert compute. Eval mode: grouped top-1 dispatch —
    each row runs exactly one (small) expert, which is the actual source
    of the per-call speedup this architecture is for.
    """
    if training:
        weights = router.gate(cond_cont, cond_cat)  # (B, n_experts)
        out = torch.zeros_like(x)
        for i, expert in enumerate(experts):
            out = out + weights[:, i : i + 1] * expert(x, cond)
        return out

    idx = router.top1(cond_cont, cond_cat)  # (B,)
    out = torch.zeros_like(x)
    for i, expert in enumerate(experts):
        mask = idx == i
        if mask.any():
            out[mask] = expert(x[mask], cond[mask])
    return out


class RoutedDenoisingMLP(nn.Module):
    """Routed drop-in for `DenoisingMLP`.

    Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all
    tiny) across experts and routes only the trunk (where the FLOPs are).
    Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as
    `DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        router: Router,
        expert_hidden_dim: int = 128,
        expert_n_blocks: int = 3,
        emb_dim: int = EMB_DIM,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        k_max: int = K_MAX,
    ) -> None:
        super().__init__()
        self.router = router
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.experts = nn.ModuleList(
            [
                ExpertTrunk(
                    x_dim,
                    expert_hidden_dim,
                    expert_n_blocks,
                    merged_cond_dim,
                    dropout=dropout,
                )
                for _ in range(router.n_experts)
            ]
        )
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, cond_out_dim),
            nn.SiLU(),
            nn.Linear(cond_out_dim, k_max + 1),
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        return _route_forward(
            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
        )

    def predict_n_sec(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)

    def pdg_embedding_weight(self) -> torch.Tensor:
        """Return the PDG embedding table weights for secondary type targets."""
        return self.cond_enc.pdg_emb.weight


class RoutedSecondaryDecoder(nn.Module):
    """Routed drop-in for `SecondaryDecoder`.

    Shares the time embedding and `SecondaryConditionEncoder` across
    experts and routes only the trunk. Same `forward` signature as
    `SecondaryDecoder`.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        router: Router,
        expert_hidden_dim: int = 128,
        expert_n_blocks: int = 3,
        emb_dim: int = EMB_DIM,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.router = router
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.experts = nn.ModuleList(
            [
                ExpertTrunk(
                    sec_dim,
                    expert_hidden_dim,
                    expert_n_blocks,
                    merged_cond_dim,
                    dropout=dropout,
                )
                for _ in range(router.n_experts)
            ]
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        return _route_forward(
            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
        )


_STAGE1_MODEL_KEYS = {
    "pdg_vocab",
    "mat_vocab",
    "hidden_dim",
    "n_blocks",
    "emb_dim",
    "dropout",
    "k_max",
}
_SEC_DECODER_MODEL_KEYS = {
    "pdg_vocab",
    "mat_vocab",
    "hidden_dim",
    "n_blocks",
    "emb_dim",
    "dropout",
}


_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")


def _parse_composed_axes(router_cfg: dict) -> list[dict]:
    """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.

    Flat keys (rather than a nested list-of-dicts) keep composed-router
    config expressible in the same one-level-of-nesting TOML/CLI shape as
    every other router option (`model.router` stays a flat table of
    scalars) — e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`,
    `axis1_type = "pdg"`, `axis1_n_experts = 3`, `axis1_emb_dim = 8`.
    Axis indices must be contiguous from 0; order follows the index, not
    dict insertion order (TOML/CLI merging doesn't preserve it reliably).
    """
    axes: dict[int, dict] = {}
    for key, value in router_cfg.items():
        m = _AXIS_KEY_RE.match(key)
        if m is None:
            continue
        idx, field = int(m.group(1)), m.group(2)
        axes.setdefault(idx, {})[field] = value
    missing = set(range(len(axes))) - axes.keys()
    if missing:
        raise ValueError(f"composed router config has gaps at axis indices {missing}")
    return [axes[i] for i in range(len(axes))]


def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router:
    """Resolve one `model.router` config into a `Router`, single-axis or composed.

    `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
    (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
    """
    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    if router_cfg["type"] == "composed":
        return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
    router_kwargs = {
        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
    }
    # Not every router needs these (EnergyRouter doesn't declare them, so
    # build_router's kwarg filtering drops them silently) but ProcessRouter
    # needs its own pdg/material embeddings sized to match the checkpoint's
    # vocab, same as the trunk's ConditionEncoder.
    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
    router_kwargs.setdefault("mat_vocab", mat_vocab)
    return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)


def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
    """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict.

    Dispatches to the routed pair when `model_config["router"]["enabled"]`
    is truthy; a missing/absent "router" key (pre-routing checkpoints)
    falls back to the monolithic pair unchanged, so this is a drop-in
    replacement for the ad-hoc constructions it replaces.
    """
    router_cfg = model_config.get("router")
    if router_cfg and router_cfg.get("enabled"):
        pdg_vocab = model_config["pdg_vocab"]
        mat_vocab = model_config["mat_vocab"]
        shared = dict(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            expert_hidden_dim=model_config.get("expert_hidden_dim", 128),
            expert_n_blocks=model_config.get("expert_n_blocks", 3),
            emb_dim=model_config.get("emb_dim", EMB_DIM),
            dropout=model_config.get("dropout", 0.1),
        )
        stage1 = RoutedDenoisingMLP(
            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
            k_max=model_config.get("k_max", K_MAX),
            **shared,
        )
        sec_decoder = RoutedSecondaryDecoder(
            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
            **shared,
        )
        return stage1, sec_decoder

    stage1 = DenoisingMLP(
        **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
    )
    sec_decoder = SecondaryDecoder(
        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
    )
    return stage1, sec_decoder