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

import torch

from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
from giant.model.network import (
    DenoisingMLP,
    EnergyRouter,
    PdgRouter,
    ProcessRouter,
    ROUTER_REGISTRY,
    RoutedDenoisingMLP,
    RoutedSecondaryDecoder,
    SecondaryDecoder,
    build_models,
    build_router,
)


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 RoutedDenoisingMLP(
        pdg_vocab=pdg,
        mat_vocab=mat,
        router=router,
        expert_hidden_dim=16,
        expert_n_blocks=2,
    )


def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
    router = build_router("energy", n_experts, **router_kwargs)
    return RoutedSecondaryDecoder(
        pdg_vocab=pdg,
        mat_vocab=mat,
        router=router,
        expert_hidden_dim=16,
        expert_n_blocks=2,
    )


# ── 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 model_config.router 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_build_router_unknown_type_raises():
    try:
        build_router("nonexistent", 4)
    except ValueError:
        return
    raise AssertionError("expected ValueError for unknown router type")


# ── 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 test_build_models_routed_with_pdg_router():
    model_config = dict(
        pdg_vocab=4,
        mat_vocab=2,
        emb_dim=16,
        dropout=0.1,
        k_max=K_MAX,
        expert_hidden_dim=16,
        expert_n_blocks=2,
        router={
            "enabled": True,
            "type": "pdg",
            "n_experts": 3,
        },
    )
    stage1, sec_decoder = build_models(model_config)
    assert isinstance(stage1, RoutedDenoisingMLP)
    assert isinstance(stage1.router, PdgRouter)
    assert len(stage1.experts) == 3
    assert stage1.router.pdg_emb.num_embeddings == 4


# ── 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():
    model_config = dict(
        pdg_vocab=4,
        mat_vocab=2,
        emb_dim=16,
        dropout=0.1,
        k_max=K_MAX,
        expert_hidden_dim=16,
        expert_n_blocks=2,
        router={
            "enabled": True,
            "type": "process",
            "n_experts": 3,
            "lambda_proc": 1.0,
        },
    )
    stage1, sec_decoder = build_models(model_config)
    assert isinstance(stage1, RoutedDenoisingMLP)
    assert isinstance(stage1.router, ProcessRouter)
    assert len(stage1.experts) == 3
    assert stage1.router.pdg_emb.num_embeddings == 4
    assert stage1.router.mat_emb.num_embeddings == 2


# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────


def test_routed_denoising_mlp_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, t, cond_cont, cond_cat)
    assert out_train.shape == (B, X_DIM)

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


def test_routed_denoising_mlp_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, t, cond_cont, cond_cat).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_denoising_mlp_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, t, cond_cont, cond_cat)

        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.router.top1(cond_cont, cond_cat)
        manual = torch.zeros_like(x_t)
        for i in range(B):
            manual[i] = model.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_denoising_mlp_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_denoising_mlp_pdg_embedding_weight_shape():
    model = _routed_stage1(pdg=5, mat=2)
    from giant.constants import EMB_DIM

    assert model.pdg_embedding_weight().shape == (5, EMB_DIM)


# ── RoutedSecondaryDecoder ───────────────────────────────────────────────────


def test_routed_secondary_decoder_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, t, cond_cont, cond_cat, stage1_out)
    assert out_train.shape == (B, SEC_DIM)

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


def test_routed_secondary_decoder_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()
    decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().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():
    model_config = dict(
        pdg_vocab=4,
        mat_vocab=2,
        hidden_dim=32,
        n_blocks=2,
        emb_dim=16,
        dropout=0.1,
        k_max=K_MAX,
    )
    stage1, sec_decoder = build_models(model_config)
    assert isinstance(stage1, DenoisingMLP)
    assert isinstance(sec_decoder, SecondaryDecoder)


def test_build_models_monolith_when_router_disabled():
    model_config = dict(
        pdg_vocab=4,
        mat_vocab=2,
        hidden_dim=32,
        n_blocks=2,
        emb_dim=16,
        dropout=0.1,
        k_max=K_MAX,
        router={"enabled": False, "type": "energy", "n_experts": 4},
    )
    stage1, sec_decoder = build_models(model_config)
    assert isinstance(stage1, DenoisingMLP)
    assert isinstance(sec_decoder, SecondaryDecoder)


def test_build_models_routed_when_enabled():
    model_config = dict(
        pdg_vocab=4,
        mat_vocab=2,
        emb_dim=16,
        dropout=0.1,
        k_max=K_MAX,
        expert_hidden_dim=16,
        expert_n_blocks=2,
        router={
            "enabled": True,
            "type": "energy",
            "n_experts": 4,
            "temperature": 0.5,
            "learn_centers": True,
            "lambda_balance": 0.0,
        },
    )
    stage1, sec_decoder = build_models(model_config)
    assert isinstance(stage1, RoutedDenoisingMLP)
    assert isinstance(sec_decoder, RoutedSecondaryDecoder)
    assert len(stage1.experts) == 4
    assert len(sec_decoder.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

    model_config = dict(
        pdg_vocab=3,
        mat_vocab=2,
        emb_dim=16,
        dropout=0.1,
        k_max=K_MAX,
        expert_hidden_dim=8,
        expert_n_blocks=1,
        router={"enabled": True, "type": "energy", "n_experts": 2},
    )
    stage1, sec_decoder = build_models(model_config)
    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)
    assert n_sec_pred.shape == (B,)

    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
        sec_decoder, 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)
