"""Frozen snapshot of `giant/model/network.py` as it stood at the v0.3.0
"step 1" commit (eb6dd27), i.e. the last commit before the step-2
composable-parts decomposition.

This is a deliberate verbatim copy, not an import of the live module — the
whole point is that this file's classes keep behaving exactly as v0.2 did
even after `giant/model/network.py` itself is rewritten, so
`tests/test_migration_v02_v03.py` has a stable "old" side to diff the new
`build_models`/`Stage1Model`/`Stage2OneShot` against (the bit-identical
acceptance test). Do not edit this file to track future
`network.py` changes — it exists specifically to stop tracking them.
"""

import inspect
import math
import re
from collections.abc import Sequence

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

from giant.constants import (
    COND_DIM,
    COND_DIM_BASE,
    EMB_DIM,
    K_MAX,
    MATERIAL_PHYS_DIM,
    PARTICLE_PHYS_DIM,
    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):
    """Fuses continuous conditioning with particle/material identity.

    Two mutually exclusive ways to turn (pdg, material) identity into the
    two `emb_dim`-wide vectors concatenated with the base continuous
    conditioning before the fusion MLP:
      - "embedding": a learned `nn.Embedding` lookup table per axis, indexed
        by `cond_cat`'s dense training-vocab index. Memorizes the training
        menu; the original Phase-2 design.
      - "physical": a small MLP per axis, mapping the axis's raw physical
        properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
        giant.data.transforms.build_features) to an `emb_dim`-wide vector —
        a drop-in replacement for the embedding lookup, computable for any
        PDG code / material name rather than only ones seen in training.
    Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the
    fusion MLP, so only how the two vectors are produced differs.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        cont_dim: int = COND_DIM,
        emb_dim: int = 16,
        out_dim: int = 128,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        if conditioning not in ("embedding", "physical"):
            raise ValueError(f"unknown conditioning mode {conditioning!r}")
        self.conditioning = conditioning
        if conditioning == "embedding":
            self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
            self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        else:
            self.particle_mlp = nn.Sequential(
                nn.Linear(PARTICLE_PHYS_DIM, emb_dim),
                nn.SiLU(),
                nn.Linear(emb_dim, emb_dim),
            )
            self.material_mlp = nn.Sequential(
                nn.Linear(MATERIAL_PHYS_DIM, emb_dim),
                nn.SiLU(),
                nn.Linear(emb_dim, emb_dim),
            )
        in_dim = COND_DIM_BASE + 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:
        if self.conditioning == "embedding":
            pdg_e = self.pdg_emb(cond_cat[:, 0])
            mat_e = self.mat_emb(cond_cat[:, 1])
        else:
            particle_phys = cond_cont[
                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
            ]
            material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
            pdg_e = self.particle_mlp(particle_phys)
            mat_e = self.material_mlp(material_phys)
        x = torch.cat([cond_cont[:, :COND_DIM_BASE], 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,
        conditioning: str = "embedding",
    ) -> 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,
            conditioning=conditioning,
        )
        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)


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,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.base = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        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, log_mass, charge) for
    one secondary ordered by descending energy — mass/charge are the
    secondary's predicted physical identity, regressed directly against real
    physics targets (see giant.data.transforms.encode_secondaries), used
    as-is with no snapping to a discrete PDG code. 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,
        conditioning: str = "embedding",
    ) -> 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,
            conditioning=conditioning,
        )
        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 WGANGenerator(nn.Module):
    """Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time.

    Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the
    input is a noise vector `z` (not a diffused/interpolated `x_t`) and the
    ResBlocks condition on the condition encoding alone (no time embedding to
    concatenate) — see `giant/model/wgan.py` for the adversarial losses, and
    `giant.sample.sample_wgan` for single-pass sampling.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        noise_dim: int = 64,
        dropout: float = 0.1,
        k_max: int = K_MAX,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.noise_dim = noise_dim
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(noise_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        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,
        z: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat)
        x = self.input_proj(z)
        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)


class Critic(nn.Module):
    """Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`.

    Kept structurally parallel to `WGANGenerator` (own condition encoder —
    separate weights from the generator's, standard GAN practice) but has no
    n_sec head: n_sec is never adversarial, it stays a plain classifier on
    the generator side.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(x_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        x: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat)
        h = self.input_proj(x)
        for block in self.blocks:
            h = block(h, cond)
        return self.out_proj(self.out_norm(h)).squeeze(-1)


class WGANSecondaryGenerator(nn.Module):
    """Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots.

    Mirrors `SecondaryDecoder` minus the time embedding, the same way
    `WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`,
    conditions on `SecondaryConditionEncoder`'s output alone.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        noise_dim: int = 64,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.noise_dim = noise_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,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(noise_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

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


class SecondaryCritic(nn.Module):
    """Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        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,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        x: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
        h = self.input_proj(x)
        for block in self.blocks:
            h = block(h, cond)
        return self.out_proj(self.out_norm(h)).squeeze(-1)


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

    def __init__(self, n_experts: int) -> None:
        super().__init__()
        self.n_experts = n_experts
        self.gumbel = False
        self.gumbel_tau = 1.0

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError

    def combine_weights(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        probs = self.gate(cond_cont, cond_cat)
        if not (self.gumbel and self.training):
            return probs
        log_probs = torch.log(probs.clamp_min(1e-8))
        return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)

    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        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 = 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:
        return torch.zeros((), device=cond_cont.device)

    def entropy_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
        return norm_entropy

    def gate_stats(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
        row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
        norm_entropy = row_entropy.mean() / math.log(self.n_experts)
        importance = gate.sum(dim=0)  # (n_experts,)
        return norm_entropy, importance


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


def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
    return lo + (hi - lo) * torch.sigmoid(raw)


def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
    p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
    return math.log(p / (1 - p))


@register_router("energy")
class EnergyRouter(Router):
    def __init__(
        self,
        n_experts: int = 4,
        temperature: float = 0.5,
        learn_centers: bool = True,
        energy_idx: int = 3,
        centers_init: Sequence[float] | None = None,
        learn_width: bool = False,
        learn_temperature: bool = False,
        width_min_ratio: float = 0.1,
        width_max_ratio: float = 10.0,
    ) -> None:
        super().__init__(n_experts)
        if learn_width and learn_temperature:
            raise ValueError("learn_width and learn_temperature are mutually exclusive")
        self.temperature = temperature
        self.energy_idx = energy_idx
        self.learn_width = learn_width
        self.learn_temperature = learn_temperature
        if learn_width or learn_temperature:
            if not (width_min_ratio < 1.0 < width_max_ratio):
                raise ValueError(
                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
                    f"({width_max_ratio}) must bracket 1.0"
                )
            self._width_lo = width_min_ratio * temperature
            self._width_hi = width_max_ratio * temperature
            raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
            if learn_width:
                self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
            else:
                self.raw_temperature = nn.Parameter(torch.tensor(raw0))
        if centers_init is None:
            centers = torch.linspace(-2.0, 2.0, n_experts)
        else:
            if len(centers_init) != n_experts:
                raise ValueError(
                    f"centers_init has {len(centers_init)} values, "
                    f"expected n_experts={n_experts}"
                )
            centers = torch.tensor(list(centers_init), dtype=torch.float32)
        if learn_centers:
            self.centers = nn.Parameter(centers)
        else:
            self.register_buffer("centers", centers)

    def effective_width(self) -> torch.Tensor | float:
        if self.learn_width:
            return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
        if self.learn_temperature:
            return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
        return self.temperature

    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.effective_width(), dim=-1)


@register_router("pdg")
class PdgRouter(Router):
    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):
    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:
        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):
    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:
        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:
    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):
    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:
    if training:
        weights = router.combine_weights(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):
    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,
        conditioning: str = "embedding",
    ) -> 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,
            conditioning=conditioning,
        )
        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:
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)


class RoutedSecondaryDecoder(nn.Module):
    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,
        conditioning: str = "embedding",
    ) -> 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,
            conditioning=conditioning,
        )
        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",
    "conditioning",
}
_SEC_DECODER_MODEL_KEYS = {
    "pdg_vocab",
    "mat_vocab",
    "hidden_dim",
    "n_blocks",
    "emb_dim",
    "dropout",
    "conditioning",
}
_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"}
_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"}
_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"}


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


def _parse_composed_axes(router_cfg: dict) -> list[dict]:
    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))]


_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")


def _check_router_conditioning_compat(
    router_types: list[str], conditioning: str
) -> None:
    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
    if bad and conditioning == "physical":
        raise ValueError(
            f"router type(s) {bad} always use a training-vocab PDG embedding, "
            "which is incompatible with conditioning='physical' (whose whole "
            "point is generalizing beyond that vocab) — pick a different "
            "router type (e.g. 'energy') or use conditioning='embedding'."
        )


def _build_router_from_cfg(
    router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
) -> Router:
    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat([a["type"] for a in axes], conditioning)
        router = build_composed_router(axes, **shared_vocab)
        router.gumbel = bool(router_cfg.get("gumbel", False))
        return router
    _check_router_conditioning_compat([router_cfg["type"]], conditioning)
    router_kwargs = {
        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
    }
    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
    router_kwargs.setdefault("mat_vocab", mat_vocab)
    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
    router.gumbel = bool(router_cfg.get("gumbel", False))
    return router


def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
    if model_config.get("mode") == "wgan":
        stage1 = WGANGenerator(
            **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
        )
        sec_decoder = WGANSecondaryGenerator(
            **{
                k: v
                for k, v in model_config.items()
                if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
            }
        )
        return stage1, sec_decoder

    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")
            or model_config.get("hidden_dim", 128),
            expert_n_blocks=model_config.get("expert_n_blocks")
            or model_config.get("n_blocks", 3),
            emb_dim=model_config.get("emb_dim", EMB_DIM),
            dropout=model_config.get("dropout", 0.1),
            conditioning=model_config.get("conditioning", "embedding"),
        )
        conditioning = shared["conditioning"]
        stage1 = RoutedDenoisingMLP(
            router=_build_router_from_cfg(
                router_cfg, pdg_vocab, mat_vocab, conditioning
            ),
            k_max=model_config.get("k_max", K_MAX),
            **shared,
        )
        sec_decoder = RoutedSecondaryDecoder(
            router=_build_router_from_cfg(
                router_cfg, pdg_vocab, mat_vocab, conditioning
            ),
            **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


def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
    critic = Critic(
        **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
    )
    sec_critic = SecondaryCritic(
        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
    )
    return critic, sec_critic