giant/tests/test_router.py

"""Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""

import pytest
import torch

from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
from giant.model.network import (
    ComposedRouter,
    EnergyRouter,
    MonolithicTrunk,
    PdgRouter,
    ProcessRouter,
    ROUTER_REGISTRY,
    RoutedTrunk,
    Stage1Model,
    Stage2OneShot,
    build_composed_router,
    build_models,
    build_router,
)

PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}


def _cond(B=8, pdg=3, mat=2):
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
    )
    return cond_cont, cond_cat


def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs):
    router = build_router("energy", n_experts, **router_kwargs)
    return Stage1Model(
        pdg_vocab=pdg,
        mat_vocab=mat,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=2,
        router=router,
        n_sec_head_k_max=K_MAX,
    )


def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
    router = build_router("energy", n_experts, **router_kwargs)
    return Stage2OneShot(
        pdg_vocab=pdg,
        mat_vocab=mat,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=2,
        generator="flow",
        time_dim=16,
        router=router,
    )


# ── Router / EnergyRouter contract ──────────────────────────────────────────


def test_energy_router_registered():
    assert ROUTER_REGISTRY["energy"] is EnergyRouter


def test_energy_router_gate_partition_of_unity():
    router = EnergyRouter(n_experts=4)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    assert g.shape == (16, 4)
    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)


def test_energy_router_top1_matches_gate_argmax():
    router = EnergyRouter(n_experts=4)
    cond_cont, cond_cat = _cond(16)
    assert torch.equal(
        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
    )


def test_energy_router_hardens_as_temperature_shrinks():
    """As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
    router = EnergyRouter(n_experts=4, temperature=1e-4)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    top1 = router.top1(cond_cont, cond_cat)
    onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)


def test_energy_router_balance_loss_is_nonnegative_scalar():
    router = EnergyRouter(n_experts=4)
    cond_cont, cond_cat = _cond(16)
    loss = router.balance_loss(cond_cont, cond_cat)
    assert loss.shape == ()
    assert loss.item() >= 0.0


def test_build_router_ignores_unrecognized_kwargs():
    # lambda_balance is a router config key but not an EnergyRouter kwarg
    router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5)
    assert isinstance(router, EnergyRouter)
    assert router.temperature == 0.3


def test_energy_router_default_centers_are_linspace():
    router = EnergyRouter(n_experts=4)
    torch.testing.assert_close(router.centers, torch.linspace(-2.0, 2.0, 4))


def test_energy_router_centers_init_overrides_default():
    centers_init = [-1.0, 0.0, 0.5, 3.0]
    router = EnergyRouter(n_experts=4, centers_init=centers_init)
    torch.testing.assert_close(router.centers, torch.tensor(centers_init))


def test_energy_router_centers_init_wrong_length_raises():
    try:
        EnergyRouter(n_experts=4, centers_init=[0.0, 1.0])
    except ValueError:
        return
    raise AssertionError("expected ValueError for centers_init length mismatch")


def test_energy_router_centers_init_respects_learn_centers_flag():
    learned = EnergyRouter(
        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
    )
    fixed = EnergyRouter(
        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
    )
    assert isinstance(learned.centers, torch.nn.Parameter)
    assert not isinstance(fixed.centers, torch.nn.Parameter)


def test_build_router_threads_centers_init_through_energy_router():
    centers_init = [-1.5, -0.5, 0.5, 1.5]
    router = build_router("energy", 4, centers_init=centers_init)
    torch.testing.assert_close(router.centers, torch.tensor(centers_init))


def test_build_router_unknown_type_raises():
    try:
        build_router("nonexistent", 4)
    except ValueError:
        return
    raise AssertionError("expected ValueError for unknown router type")


# ── EnergyRouter learn_width / learn_temperature ────────────────────────────


def test_energy_router_learn_width_matches_fixed_temperature_at_init():
    """Enabling learn_width should be a no-op at init — the warm-started
    per-expert width must reproduce the fixed-temperature gate exactly."""
    centers_init = [-1.0, 0.0, 0.5, 1.5]
    fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
    learned = EnergyRouter(
        n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
    )
    cond_cont, cond_cat = _cond(16)
    torch.testing.assert_close(
        learned.gate(cond_cont, cond_cat),
        fixed.gate(cond_cont, cond_cat),
        atol=1e-5,
        rtol=0,
    )


def test_energy_router_learn_temperature_matches_fixed_temperature_at_init():
    centers_init = [-1.0, 0.0, 0.5, 1.5]
    fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
    learned = EnergyRouter(
        n_experts=4,
        temperature=0.3,
        centers_init=centers_init,
        learn_temperature=True,
    )
    cond_cont, cond_cat = _cond(16)
    torch.testing.assert_close(
        learned.gate(cond_cont, cond_cat),
        fixed.gate(cond_cont, cond_cat),
        atol=1e-5,
        rtol=0,
    )


def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
    try:
        EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
    except ValueError:
        return
    raise AssertionError(
        "expected ValueError for learn_width and learn_temperature both set"
    )


def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
    try:
        EnergyRouter(
            n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
        )
    except ValueError:
        return
    raise AssertionError(
        "expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
    )


def test_energy_router_effective_width_stays_within_bounds():
    router = EnergyRouter(
        n_experts=4,
        temperature=0.5,
        learn_width=True,
        width_min_ratio=0.1,
        width_max_ratio=10.0,
    )
    lo, hi = 0.1 * 0.5, 10.0 * 0.5
    with torch.no_grad():
        router.raw_width.fill_(1e6)
    width = router.effective_width()
    assert isinstance(width, torch.Tensor)
    assert torch.all(width <= hi + 1e-4)
    with torch.no_grad():
        router.raw_width.fill_(-1e6)
    width = router.effective_width()
    assert isinstance(width, torch.Tensor)
    assert torch.all(width >= lo - 1e-4)


def test_energy_router_learn_width_gate_still_partition_of_unity():
    router = EnergyRouter(n_experts=4, learn_width=True)
    with torch.no_grad():
        router.raw_width.copy_(torch.randn(4) * 3)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)


def test_energy_router_learn_width_hardens_when_pushed_to_floor():
    """Pushing every expert's width toward the (tiny) floor should harden the
    gate to a one-hot at the nearest center, generalizing the fixed-
    temperature->0 hardening test to the per-expert path."""
    router = EnergyRouter(
        n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
    )
    with torch.no_grad():
        router.raw_width.fill_(-1e6)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    e = cond_cont[:, router.energy_idx].unsqueeze(-1)
    d2 = (e - router.centers.unsqueeze(0)) ** 2
    onehot = torch.nn.functional.one_hot(d2.argmin(dim=-1), num_classes=4).float()
    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)


def test_energy_router_own_width_controls_own_coverage_independent_of_others():
    """Widening one expert's width should monotonically grow only that
    expert's own gate share, without needing to touch any other expert's
    width — the "each expert learns its own coverage independently" property
    this feature is meant to add."""
    router = EnergyRouter(
        n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
    )
    cond_cont, cond_cat = _cond(4)
    cond_cont[:, 3] = 3.0  # fixed energy, unequal distance to each center

    shares = []
    with torch.no_grad():
        for raw in torch.linspace(-8.0, 8.0, 9):
            router.raw_width[0] = raw
            shares.append(router.gate(cond_cont, cond_cat)[0, 0].item())

    assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:]))


def test_build_router_threads_learn_width_kwargs_through():
    router = build_router(
        "energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
    )
    assert isinstance(router, EnergyRouter)
    assert router.learn_width is True
    assert isinstance(router.raw_width, torch.nn.Parameter)
    assert router.raw_width.shape == (4,)


def test_router_entropy_loss_is_nonnegative_bounded_scalar():
    router = EnergyRouter(n_experts=4)
    cond_cont, cond_cat = _cond(16)
    loss = router.entropy_loss(cond_cont, cond_cat)
    assert loss.shape == ()
    assert 0.0 <= loss.item() <= 1.0


# ── Router.combine_weights (straight-through Gumbel-softmax) ───────────────


def test_combine_weights_defaults_to_gate():
    """gumbel=False (the default) must be a pure pass-through to gate()."""
    router = EnergyRouter(n_experts=4)
    cond_cont, cond_cat = _cond(16)
    torch.testing.assert_close(
        router.combine_weights(cond_cont, cond_cat),
        router.gate(cond_cont, cond_cat),
    )


def test_combine_weights_gumbel_train_mode_is_hard_one_hot():
    router = EnergyRouter(n_experts=4)
    router.gumbel = True
    router.gumbel_tau = 0.5
    router.train()
    cond_cont, cond_cat = _cond(16)
    weights = router.combine_weights(cond_cont, cond_cat)
    assert weights.shape == (16, 4)
    torch.testing.assert_close(weights.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
    assert torch.all((weights.max(dim=-1).values - 1.0).abs() < 1e-5)


def test_combine_weights_gumbel_eval_mode_falls_back_to_gate():
    """No Gumbel noise at eval — combine_weights must match gate() exactly,
    same as the gumbel=False path, once the router is in eval mode."""
    router = EnergyRouter(n_experts=4)
    router.gumbel = True
    router.eval()
    cond_cont, cond_cat = _cond(16)
    torch.testing.assert_close(
        router.combine_weights(cond_cont, cond_cat),
        router.gate(cond_cont, cond_cat),
    )


def test_combine_weights_gumbel_straight_through_gradient_reaches_centers():
    router = EnergyRouter(n_experts=4, learn_centers=True)
    router.gumbel = True
    router.gumbel_tau = 0.5
    router.train()
    cond_cont, cond_cat = _cond(16)
    weights = router.combine_weights(cond_cont, cond_cat)
    weights.sum().backward()
    assert router.centers.grad is not None
    assert torch.any(router.centers.grad != 0.0)


def test_build_router_from_cfg_sets_gumbel_from_config():
    from giant.model.network import _build_router_from_cfg

    router = _build_router_from_cfg(
        {"enabled": True, "type": "energy", "n_experts": 4, "gumbel": True},
        pdg_vocab=3,
        mat_vocab=2,
    )
    assert router.gumbel is True

    router_off = _build_router_from_cfg(
        {"enabled": True, "type": "energy", "n_experts": 4},
        pdg_vocab=3,
        mat_vocab=2,
    )
    assert router_off.gumbel is False


def test_build_router_from_cfg_sets_gumbel_for_composed_router():
    from giant.model.network import _build_router_from_cfg

    router = _build_router_from_cfg(
        {
            "enabled": True,
            "type": "composed",
            "gumbel": True,
            "axis0_type": "energy",
            "axis0_n_experts": 4,
            "axis1_type": "pdg",
            "axis1_n_experts": 3,
        },
        pdg_vocab=5,
        mat_vocab=2,
    )
    assert isinstance(router, ComposedRouter)
    assert router.gumbel is True


def test_routed_stage1_forward_runs_with_gumbel_enabled():
    """End-to-end forward through _route_forward's train branch with
    straight-through Gumbel-softmax combine weights enabled."""
    B = 8
    model = _routed_stage1(n_experts=3)
    model.trunk.router.gumbel = True
    model.trunk.router.gumbel_tau = 0.5
    model.train()
    x_t = torch.randn(B, X_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    out = model(x_t, cond_cont, cond_cat, t=t)
    assert out.shape == (B, X_DIM)
    assert torch.isfinite(out).all()


# ── PdgRouter ────────────────────────────────────────────────────────────────


def test_pdg_router_registered():
    assert ROUTER_REGISTRY["pdg"] is PdgRouter


def test_pdg_router_gate_partition_of_unity():
    router = PdgRouter(n_experts=4, pdg_vocab=3)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    assert g.shape == (16, 4)
    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)


def test_pdg_router_top1_matches_gate_argmax():
    router = PdgRouter(n_experts=4, pdg_vocab=3)
    cond_cont, cond_cat = _cond(16)
    assert torch.equal(
        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
    )


def test_pdg_router_hardens_as_temperature_shrinks():
    """As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
    router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    top1 = router.top1(cond_cont, cond_cat)
    onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)


def test_pdg_router_balance_loss_is_nonnegative_scalar():
    router = PdgRouter(n_experts=4, pdg_vocab=3)
    cond_cont, cond_cat = _cond(16)
    loss = router.balance_loss(cond_cont, cond_cat)
    assert loss.shape == ()
    assert loss.item() >= 0.0


def test_pdg_router_classify_loss_defaults_to_zero():
    """PDG is already known at gate time (unlike ProcessRouter's process
    label), so no supervision is needed — falls back to Router's default."""
    router = PdgRouter(n_experts=4, pdg_vocab=3)
    cond_cont, cond_cat = _cond(16)
    labels = torch.randint(0, 4, (16,))
    loss = router.classify_loss(cond_cont, cond_cat, labels)
    assert loss.shape == ()
    assert loss.item() == 0.0


def test_pdg_router_only_reads_pdg_column():
    """Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material."""
    router = PdgRouter(n_experts=4, pdg_vocab=3)
    cond_cont, cond_cat = _cond(16)
    g_before = router.gate(cond_cont, cond_cat)

    cond_cont_perturbed = torch.randn_like(cond_cont)
    cond_cat_diff_mat = cond_cat.clone()
    cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2
    g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat)

    torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0)


def test_build_router_pdg_type_uses_pdg_vocab():
    router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
    assert isinstance(router, PdgRouter)
    assert router.pdg_emb.num_embeddings == 5


def _nested_cfg(
    pdg_vocab,
    mat_vocab,
    stage1_router=None,
    stage2_router=None,
    particle_type="physical",
    material_type="physical",
    **overrides,
):
    """Minimal new-shape (v0.3.0) model_config for build_models, with
    optional router sub-blocks. `overrides` deep-patches stage1_model."""
    stage1_model = {
        "active": True,
        "generator": "flow",
        "hidden_dim": 16,
        "n_res_blocks": 2,
        "dropout": 0.0,
        "flow": {"time_dim": 16},
        "router": stage1_router or {"enabled": False},
    }
    stage1_model.update(overrides)
    return {
        "pdg_vocab": pdg_vocab,
        "mat_vocab": mat_vocab,
        "conditioning": {
            "out_dim": 32,
            "particle": {"type": particle_type, "emb_dim": 8, "n_layers": 1},
            "material": {"type": material_type, "emb_dim": 8, "n_layers": 1},
        },
        "stage1_model": stage1_model,
        "stage2_model": {
            "active": True,
            "decoder": "one_shot",
            "generator": "flow",
            "hidden_dim": 16,
            "n_res_blocks": 2,
            "dropout": 0.0,
            "k_max": K_MAX,
            "context_dim": 16,
            "n_sec": {"mode": "head"},
            "flow": {"time_dim": 16},
            "router": stage2_router or {"enabled": False, "tie_to_stage1": False},
        },
    }


def test_build_models_routed_with_pdg_router():
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        particle_type="embedding",
        material_type="embedding",
        stage1_router={"enabled": True, "type": "pdg", "n_experts": 3},
    )
    models = build_models(cfg)
    stage1 = models["stage1"]
    assert isinstance(stage1, Stage1Model)
    assert isinstance(stage1.trunk, RoutedTrunk)
    assert isinstance(stage1.trunk.router, PdgRouter)
    assert len(stage1.trunk.experts) == 3
    assert stage1.trunk.router.pdg_emb.num_embeddings == 4


def test_build_models_rejects_pdg_router_with_physical_conditioning():
    """conditioning.particle.type="physical" is meant to generalize beyond the
    training PDG vocab; PdgRouter always uses a training-vocab nn.Embedding
    regardless of conditioning, so the combination must raise rather than
    silently building a model that can't actually generalize the way it
    claims to."""
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        stage1_router={"enabled": True, "type": "pdg", "n_experts": 3},
    )
    with pytest.raises(ValueError, match="physical"):
        build_models(cfg)


def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning():
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        stage1_router={
            "enabled": True,
            "type": "composed",
            "axis0_type": "energy",
            "axis0_n_experts": 2,
            "axis1_type": "pdg",
            "axis1_n_experts": 3,
        },
    )
    with pytest.raises(ValueError, match="physical"):
        build_models(cfg)


# ── ProcessRouter ────────────────────────────────────────────────────────────


def test_process_router_registered():
    assert ROUTER_REGISTRY["process"] is ProcessRouter


def test_process_router_gate_partition_of_unity():
    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
    cond_cont, cond_cat = _cond(16)
    g = router.gate(cond_cont, cond_cat)
    assert g.shape == (16, 4)
    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)


def test_process_router_top1_matches_gate_argmax():
    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
    cond_cont, cond_cat = _cond(16)
    assert torch.equal(
        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
    )


def test_process_router_balance_loss_is_nonnegative_scalar():
    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
    cond_cont, cond_cat = _cond(16)
    loss = router.balance_loss(cond_cont, cond_cat)
    assert loss.shape == ()
    assert loss.item() >= 0.0


def test_process_router_classify_loss_decreases_with_training():
    """The classifier should be able to fit an arbitrary label assignment —
    a sanity check that gradients actually flow to the process classifier."""
    torch.manual_seed(0)
    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
    cond_cont, cond_cat = _cond(32)
    labels = torch.randint(0, 4, (32,))

    opt = torch.optim.Adam(router.parameters(), lr=0.05)
    first = router.classify_loss(cond_cont, cond_cat, labels).item()
    for _ in range(50):
        opt.zero_grad()
        loss = router.classify_loss(cond_cont, cond_cat, labels)
        loss.backward()
        opt.step()
    last = loss.item()
    assert last < first


def test_energy_router_classify_loss_defaults_to_zero():
    """Routers with no supervised signal (EnergyRouter) fall back to the
    Router base class's zero-loss default."""
    router = EnergyRouter(n_experts=4)
    cond_cont, cond_cat = _cond(16)
    labels = torch.randint(0, 4, (16,))
    loss = router.classify_loss(cond_cont, cond_cat, labels)
    assert loss.shape == ()
    assert loss.item() == 0.0


def test_build_router_process_type_uses_pdg_mat_vocab():
    router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
    assert isinstance(router, ProcessRouter)
    assert router.pdg_emb.num_embeddings == 5
    assert router.mat_emb.num_embeddings == 3


def test_build_models_routed_with_process_router():
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        particle_type="embedding",
        material_type="embedding",
        stage1_router={
            "enabled": True,
            "type": "process",
            "n_experts": 3,
            "lambda_proc": 1.0,
        },
    )
    models = build_models(cfg)
    stage1 = models["stage1"]
    assert isinstance(stage1, Stage1Model)
    assert isinstance(stage1.trunk, RoutedTrunk)
    assert isinstance(stage1.trunk.router, ProcessRouter)
    assert len(stage1.trunk.experts) == 3
    assert stage1.trunk.router.pdg_emb.num_embeddings == 4
    assert stage1.trunk.router.mat_emb.num_embeddings == 2


# ── ComposedRouter ───────────────────────────────────────────────────────────


def test_composed_router_n_experts_is_product():
    router = ComposedRouter(
        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
    )
    assert router.n_experts == 12


def test_composed_router_gate_partition_of_unity():
    router = ComposedRouter(
        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
    )
    cond_cont, cond_cat = _cond(16, pdg=5)
    g = router.gate(cond_cont, cond_cat)
    assert g.shape == (16, 12)
    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)


def test_composed_router_gate_is_outer_product_of_sub_gates():
    energy_router = EnergyRouter(n_experts=4)
    pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
    router = ComposedRouter([energy_router, pdg_router])
    cond_cont, cond_cat = _cond(16, pdg=5)

    g_energy = energy_router.gate(cond_cont, cond_cat)  # (16, 4)
    g_pdg = pdg_router.gate(cond_cont, cond_cat)  # (16, 3)
    expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1)  # (16, 12)

    torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected)


def test_composed_router_top1_factors_into_per_axis_argmax():
    """Joint argmax over the outer product must equal the pair of per-axis
    argmaxes, flattened with the same row-major index convention as gate()."""
    energy_router = EnergyRouter(n_experts=4)
    pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
    router = ComposedRouter([energy_router, pdg_router])
    cond_cont, cond_cat = _cond(16, pdg=5)

    joint_idx = router.top1(cond_cont, cond_cat)
    energy_idx = energy_router.top1(cond_cont, cond_cat)
    pdg_idx = pdg_router.top1(cond_cont, cond_cat)
    expected = energy_idx * pdg_router.n_experts + pdg_idx

    assert torch.equal(joint_idx, expected)


def test_composed_router_supports_different_expert_counts_per_axis():
    router = ComposedRouter(
        [EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
    )
    assert router.n_experts == 10
    cond_cont, cond_cat = _cond(8, pdg=5)
    assert router.gate(cond_cont, cond_cat).shape == (8, 10)


def test_composed_router_classify_loss_sums_sub_router_losses():
    """energy/pdg both default to zero, so the composed loss should too."""
    router = ComposedRouter(
        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
    )
    cond_cont, cond_cat = _cond(16, pdg=5)
    labels = torch.randint(0, 4, (16,))
    loss = router.classify_loss(cond_cont, cond_cat, labels)
    assert loss.shape == ()
    assert loss.item() == 0.0


def test_composed_router_rejects_empty_router_list():
    try:
        ComposedRouter([])
    except ValueError:
        return
    raise AssertionError("expected ValueError for empty router list")


def test_composed_router_not_in_registry():
    assert "composed" not in ROUTER_REGISTRY


# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ───────


def test_parse_composed_axes_groups_indexed_keys():
    from giant.model.network import _parse_composed_axes

    router_cfg = {
        "enabled": True,
        "type": "composed",
        "axis0_type": "energy",
        "axis0_n_experts": 4,
        "axis0_temperature": 0.3,
        "axis1_type": "pdg",
        "axis1_n_experts": 3,
        "axis1_emb_dim": 6,
    }
    axes = _parse_composed_axes(router_cfg)
    assert axes == [
        {"type": "energy", "n_experts": 4, "temperature": 0.3},
        {"type": "pdg", "n_experts": 3, "emb_dim": 6},
    ]


def test_parse_composed_axes_ignores_unrelated_keys():
    from giant.model.network import _parse_composed_axes

    router_cfg = {
        "enabled": True,
        "type": "composed",
        "lambda_balance": 0.0,
        "axis0_type": "energy",
        "axis0_n_experts": 4,
    }
    axes = _parse_composed_axes(router_cfg)
    assert axes == [{"type": "energy", "n_experts": 4}]


def test_parse_composed_axes_raises_on_index_gap():
    from giant.model.network import _parse_composed_axes

    router_cfg = {
        "type": "composed",
        "axis0_type": "energy",
        "axis0_n_experts": 4,
        # axis1 missing entirely
        "axis2_type": "pdg",
        "axis2_n_experts": 3,
    }
    try:
        _parse_composed_axes(router_cfg)
    except ValueError:
        return
    raise AssertionError("expected ValueError for a gap in axis indices")


def test_build_composed_router_resolves_per_axis_specs():
    router = build_composed_router(
        [
            {"type": "energy", "n_experts": 4, "temperature": 0.3},
            {"type": "pdg", "n_experts": 3, "emb_dim": 6},
        ],
        pdg_vocab=5,
        mat_vocab=2,
    )
    assert isinstance(router, ComposedRouter)
    assert router.n_experts == 12
    energy_router, pdg_router = router.routers
    assert isinstance(energy_router, EnergyRouter)
    assert energy_router.temperature == 0.3
    assert isinstance(pdg_router, PdgRouter)
    assert pdg_router.pdg_emb.num_embeddings == 5
    assert pdg_router.pdg_emb.embedding_dim == 6


def test_build_models_routed_with_composed_router():
    cfg = _nested_cfg(
        pdg_vocab=5,
        mat_vocab=2,
        particle_type="embedding",
        material_type="embedding",
        stage1_router={
            "enabled": True,
            "type": "composed",
            "axis0_type": "energy",
            "axis0_n_experts": 4,
            "axis1_type": "pdg",
            "axis1_n_experts": 3,
        },
        stage2_router={
            "enabled": True,
            "tie_to_stage1": False,
            "type": "composed",
            "axis0_type": "energy",
            "axis0_n_experts": 4,
            "axis1_type": "pdg",
            "axis1_n_experts": 3,
        },
    )
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert stage1 is not None and stage2 is not None
    assert isinstance(stage1.trunk, RoutedTrunk)
    assert isinstance(stage1.trunk.router, ComposedRouter)
    assert len(stage1.trunk.experts) == 12
    assert len(stage2.trunk.experts) == 12
    # stage1 and stage2 must not share router weights when tie_to_stage1 is
    # false (same convention as v0.2's two-independent-routers behaviour).
    assert stage1.trunk.router is not stage2.trunk.router


def test_build_models_routed_stage2_ties_to_stage1_router():
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        stage1_router={"enabled": True, "type": "energy", "n_experts": 3},
        stage2_router={
            "enabled": True,
            "tie_to_stage1": True,
            "type": "energy",
            "n_experts": 3,
        },
    )
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert stage1 is not None and stage2 is not None
    assert stage1.trunk.router is stage2.trunk.router


def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
    from giant.sample import sample_flow, sample_secondaries

    cfg = _nested_cfg(
        pdg_vocab=3,
        mat_vocab=2,
        # PdgRouter (axis1) always builds its own training-vocab embedding,
        # incompatible with conditioning.particle.type="physical" (the
        # _nested_cfg default) — see _check_router_conditioning_compat.
        particle_type="embedding",
        material_type="embedding",
        stage1_router={
            "enabled": True,
            "type": "composed",
            "axis0_type": "energy",
            "axis0_n_experts": 2,
            "axis1_type": "pdg",
            "axis1_n_experts": 2,
        },
    )
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert stage1 is not None and stage2 is not None
    B = 5
    cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
    stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
    assert stage1_norm.shape == (B, X_DIM)
    # A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) —
    # sample_flow returns n_sec_pred=None here, and n_sec must be
    # asked of stage2 instead, using the just-sampled stage1_norm as context.
    assert n_sec_pred is None
    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
    assert n_sec_pred.shape == (B,)

    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
    )
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_valid.shape == (B, K_MAX)


# ── Stage1Model with a routed trunk ─────────────────────────────────────────


def test_routed_stage1_output_shape_train_and_eval():
    B = 8
    model = _routed_stage1()
    x_t = torch.randn(B, X_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)

    model.train()
    out_train = model(x_t, cond_cont, cond_cat, t=t)
    assert out_train.shape == (B, X_DIM)

    model.eval()
    with torch.no_grad():
        out_eval = model(x_t, cond_cont, cond_cat, t=t)
    assert out_eval.shape == (B, X_DIM)


def test_routed_stage1_gradients_flow_in_train_mode():
    """Soft mixture in train mode should touch every expert's parameters."""
    B = 8
    model = _routed_stage1(n_experts=3)
    x_t = torch.randn(B, X_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    model.train()
    flow_loss = model(x_t, cond_cont, cond_cat, t=t).sum()
    nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
    (flow_loss + nsec_loss).backward()
    for name, p in model.named_parameters():
        assert p.grad is not None, f"no grad for {name}"


def test_routed_stage1_eval_dispatch_matches_manual_grouping():
    """Eval-mode grouped top-1 dispatch must equal running each row through
    its assigned expert individually (batch order shouldn't matter)."""
    B = 12
    model = _routed_stage1(n_experts=4)
    model.eval()
    x_t = torch.randn(B, X_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)

    with torch.no_grad():
        batched = model(x_t, cond_cont, cond_cat, t=t)

        t_emb = model.time_emb(t)
        c_emb = model.cond_enc(cond_cont, cond_cat)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        idx = model.trunk.router.top1(cond_cont, cond_cat)
        manual = torch.zeros_like(x_t)
        for i in range(B):
            manual[i] = model.trunk.experts[int(idx[i])](
                x_t[i : i + 1], cond[i : i + 1]
            )[0]

    torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)


def test_routed_stage1_predict_n_sec_shape():
    B = 6
    model = _routed_stage1()
    cond_cont, cond_cat = _cond(B)
    logits = model.predict_n_sec(cond_cont, cond_cat)
    assert logits.shape == (B, K_MAX + 1)


def test_routed_stage1_has_no_pdg_embedding_weight_method():
    model = _routed_stage1(pdg=5, mat=2)
    assert not hasattr(model, "pdg_embedding_weight")


# ── Stage2OneShot with a routed trunk ────────────────────────────────────────


def test_routed_stage2_output_shape_train_and_eval():
    B = 8
    decoder = _routed_sec_decoder()
    x_t = torch.randn(B, SEC_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)

    decoder.train()
    out_train = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
    assert out_train.shape == (B, SEC_DIM)

    decoder.eval()
    with torch.no_grad():
        out_eval = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
    assert out_eval.shape == (B, SEC_DIM)


def test_routed_stage2_gradients_flow():
    B = 4
    decoder = _routed_sec_decoder(n_experts=3)
    x_t = torch.randn(B, SEC_DIM)
    t = torch.rand(B)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    decoder.train()
    flow_loss = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t).sum()
    nsec_loss = decoder.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
    (flow_loss + nsec_loss).backward()
    for name, p in decoder.named_parameters():
        assert p.grad is not None, f"no grad for {name}"


# ── build_models dispatch ────────────────────────────────────────────────────


def test_build_models_monolith_when_router_absent():
    cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2)
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert isinstance(stage1, Stage1Model)
    assert isinstance(stage2, Stage2OneShot)
    assert isinstance(stage1.trunk, MonolithicTrunk)
    assert isinstance(stage2.trunk, MonolithicTrunk)


def test_build_models_monolith_when_router_disabled():
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        stage1_router={"enabled": False, "type": "energy", "n_experts": 4},
    )
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert stage1 is not None and stage2 is not None
    assert isinstance(stage1.trunk, MonolithicTrunk)
    assert isinstance(stage2.trunk, MonolithicTrunk)


def test_build_models_routed_when_enabled():
    cfg = _nested_cfg(
        pdg_vocab=4,
        mat_vocab=2,
        stage1_router={
            "enabled": True,
            "type": "energy",
            "n_experts": 4,
            "temperature": 0.5,
            "learn_centers": True,
            "lambda_balance": 0.0,
        },
        stage2_router={
            "enabled": True,
            "tie_to_stage1": False,
            "type": "energy",
            "n_experts": 4,
            "temperature": 0.5,
            "learn_centers": True,
        },
    )
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert stage1 is not None and stage2 is not None
    assert isinstance(stage1.trunk, RoutedTrunk)
    assert isinstance(stage2.trunk, RoutedTrunk)
    assert len(stage1.trunk.experts) == 4
    assert len(stage2.trunk.experts) == 4


def test_build_models_routed_pair_is_drop_in_for_sample_flow():
    """Exercise the exact calling convention giant/sample.py uses."""
    from giant.sample import sample_flow, sample_secondaries

    cfg = _nested_cfg(
        pdg_vocab=3,
        mat_vocab=2,
        stage1_router={"enabled": True, "type": "energy", "n_experts": 2},
    )
    models = build_models(cfg)
    stage1, stage2 = models["stage1"], models["stage2"]
    assert stage1 is not None and stage2 is not None
    B = 5
    cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
    stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
    assert stage1_norm.shape == (B, X_DIM)
    # A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) —
    # sample_flow returns n_sec_pred=None here, and n_sec must be
    # asked of stage2 instead, using the just-sampled stage1_norm as context.
    assert n_sec_pred is None
    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
    assert n_sec_pred.shape == (B,)

    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
    )
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_valid.shape == (B, K_MAX)