"""Migration acceptance test for v0.3.0 step 2: "load a v0.2 checkpoint
through migrate_config + the new build_models, and diff its outputs against
v0.2 code on the same input batch — bit-identical, or the refactor has
changed something it should not have."

No `/ceph` access on this machine (see CLAUDE.md's Compute environment
section), so a real trained checkpoint can't be used here — a separate
portal-machine follow-up with a real checkpoint is planned instead. This
test is the synthetic stand-in: build a v0.2-shaped
model from the frozen `tests/legacy/network_v02_snapshot.py` classes with
fixed-seed random weights (playing the role of "a v0.2 checkpoint"), migrate
its config and remap its state dict onto the new `build_models` output, and
assert the two produce bit-identical output on the same random input batch.
"""

import torch

from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
from giant.model import network as net
from tests.legacy import network_v02_snapshot as legacy

PDG_VOCAB = 12
MAT_VOCAB = 4
HIDDEN_DIM = 32
N_BLOCKS = 2
EMB_DIM = 8
K = 6  # small k_max for a fast test
BATCH = 5


def _legacy_model_config(mode: str, conditioning: str) -> dict:
    return {
        "pdg_vocab": PDG_VOCAB,
        "mat_vocab": MAT_VOCAB,
        "hidden_dim": HIDDEN_DIM,
        "n_blocks": N_BLOCKS,
        "emb_dim": EMB_DIM,
        "dropout": 0.0,
        "k_max": K,
        "conditioning": conditioning,
        "router": {"enabled": False},
        "mode": mode,
        "noise_dim": 16,
    }


def _random_batch(seed: int):
    g = torch.Generator().manual_seed(seed)
    cond_cont = torch.randn(BATCH, COND_DIM, generator=g)
    cond_cat = torch.randint(0, min(PDG_VOCAB, MAT_VOCAB), (BATCH, 2), generator=g)
    x1 = torch.randn(BATCH, X_DIM, generator=g)
    x2 = torch.randn(BATCH, K * SEC_SLOT_DIM, generator=g)
    t = torch.rand(BATCH, generator=g)
    return cond_cont, cond_cat, x1, x2, t


def _assert_bit_identical(a: torch.Tensor, b: torch.Tensor, label: str) -> None:
    assert a.shape == b.shape, f"{label}: shape mismatch {a.shape} vs {b.shape}"
    assert torch.equal(a, b), f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"


def _run_migration_check(mode: str, conditioning: str) -> None:
    torch.manual_seed(0)
    legacy_cfg = _legacy_model_config(mode, conditioning)

    if mode == "wgan":
        old_stage1 = legacy.WGANGenerator(
            pdg_vocab=PDG_VOCAB,
            mat_vocab=MAT_VOCAB,
            hidden_dim=HIDDEN_DIM,
            n_blocks=N_BLOCKS,
            emb_dim=EMB_DIM,
            noise_dim=16,
            dropout=0.0,
            k_max=K,
            conditioning=conditioning,
        )
        old_stage2 = legacy.WGANSecondaryGenerator(
            pdg_vocab=PDG_VOCAB,
            mat_vocab=MAT_VOCAB,
            hidden_dim=HIDDEN_DIM,
            n_blocks=N_BLOCKS,
            emb_dim=EMB_DIM,
            sec_dim=K * SEC_SLOT_DIM,
            noise_dim=16,
            dropout=0.0,
            conditioning=conditioning,
        )
    else:
        old_stage1 = legacy.DenoisingMLP(
            pdg_vocab=PDG_VOCAB,
            mat_vocab=MAT_VOCAB,
            hidden_dim=HIDDEN_DIM,
            n_blocks=N_BLOCKS,
            emb_dim=EMB_DIM,
            dropout=0.0,
            k_max=K,
            conditioning=conditioning,
        )
        old_stage2 = legacy.SecondaryDecoder(
            pdg_vocab=PDG_VOCAB,
            mat_vocab=MAT_VOCAB,
            hidden_dim=HIDDEN_DIM,
            n_blocks=N_BLOCKS,
            emb_dim=EMB_DIM,
            sec_dim=K * SEC_SLOT_DIM,
            dropout=0.0,
            conditioning=conditioning,
        )
    old_stage1.eval()
    old_stage2.eval()

    cond_cont, cond_cat, x1, x2, t = _random_batch(seed=123)
    z1 = torch.randn(BATCH, 16, generator=torch.Generator().manual_seed(456))
    z2 = torch.randn(BATCH, 16, generator=torch.Generator().manual_seed(789))

    with torch.no_grad():
        if mode == "wgan":
            old_out1 = old_stage1(z1, cond_cont, cond_cat)
        else:
            old_out1 = old_stage1(x1, t, cond_cont, cond_cat)
        old_n_sec = old_stage1.predict_n_sec(cond_cont, cond_cat)
        if mode == "wgan":
            old_out2 = old_stage2(z2, cond_cont, cond_cat, old_out1)
        else:
            old_out2 = old_stage2(x2, t, cond_cont, cond_cat, old_out1)

    # --- migrate: config + state dict, through the new build_models ---
    new_models = net.build_models(legacy_cfg)
    new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"]
    assert isinstance(new_stage1, net.Stage1Model)
    assert isinstance(new_stage2, net.Stage2OneShot)
    # n_sec.owner="stage1": n_sec lives on stage1, not stage2, for a
    # migrated v0.2 checkpoint.
    assert new_stage1.n_sec_head is not None
    assert new_stage2.n_sec_head is None

    remapped1, remapped2 = net.migrate_legacy_state_dict(old_stage1.state_dict(), old_stage2.state_dict())
    missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
    missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
    assert not missing1 and not unexpected1
    assert not missing2 and not unexpected2
    new_stage1.eval()
    new_stage2.eval()

    with torch.no_grad():
        if mode == "wgan":
            new_out1 = new_stage1(z1, cond_cont, cond_cat)
        else:
            new_out1 = new_stage1(x1, cond_cont, cond_cat, t=t)
        new_n_sec = new_stage1.predict_n_sec(cond_cont, cond_cat)
        if mode == "wgan":
            new_out2 = new_stage2(z2, cond_cont, cond_cat, new_out1)
        else:
            new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)

    _assert_bit_identical(old_out1, new_out1, f"stage1 output ({mode}, {conditioning})")
    _assert_bit_identical(old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})")
    _assert_bit_identical(old_out2, new_out2, f"stage2 output ({mode}, {conditioning})")


def test_migration_flow_embedding():
    _run_migration_check(mode="flow", conditioning="embedding")


def test_migration_flow_physical():
    _run_migration_check(mode="flow", conditioning="physical")


def test_migration_wgan_embedding():
    _run_migration_check(mode="wgan", conditioning="embedding")


def test_migration_wgan_physical():
    _run_migration_check(mode="wgan", conditioning="physical")


def test_migrate_legacy_model_config_shape():
    """_migrate_legacy_model_config produces the nested shape build_models
    expects, with the n_sec.owner marker set so build_models routes the
    n_sec head back onto stage 1."""
    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
    migrated = net._migrate_legacy_model_config(legacy_cfg)
    assert migrated["pdg_vocab"] == PDG_VOCAB
    assert migrated["mat_vocab"] == MAT_VOCAB
    assert migrated["conditioning"]["particle"]["type"] == "physical"
    assert migrated["conditioning"]["particle"]["n_layers"] == 2
    assert migrated["conditioning"]["material"]["n_layers"] == 2
    assert migrated["stage1_model"]["hidden_dim"] == HIDDEN_DIM
    assert migrated["stage2_model"]["n_sec"]["owner"] == "stage1"
    assert migrated["stage2_model"]["decoder"] == "one_shot"


def test_migrate_legacy_model_config_nonzero_expert_dims_raises():
    """Regression: a v0.2 checkpoint's model_config carrying a non-default
    expert_hidden_dim/expert_n_blocks must fail loudly through this path too
    — not just giant.config.migrate_config's parallel TOML-load path.
    Silently dropping these keys (build_router's kwarg filtering) would
    resize the experts instead of refusing."""
    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
    legacy_cfg["router"] = {
        "enabled": True,
        "expert_hidden_dim": 128,
        "expert_n_blocks": 0,
    }
    try:
        net._migrate_legacy_model_config(legacy_cfg)
        assert False, "expected ValueError"
    except ValueError as e:
        assert "expert_hidden_dim" in str(e)


def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently():
    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
    legacy_cfg["router"] = {
        "enabled": True,
        "expert_hidden_dim": 0,
        "expert_n_blocks": 0,
    }
    migrated = net._migrate_legacy_model_config(legacy_cfg)
    assert "expert_hidden_dim" not in migrated["stage1_model"]["router"]
    assert "expert_n_blocks" not in migrated["stage1_model"]["router"]
    assert "expert_hidden_dim" not in migrated["stage2_model"]["router"]
    assert "expert_n_blocks" not in migrated["stage2_model"]["router"]


def test_build_models_with_legacy_config_nonzero_expert_dims_raises():
    """The same check must also fire through the actual caller,
    build_models, not just the internal helper directly."""
    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
    legacy_cfg["router"] = {
        "enabled": True,
        "expert_hidden_dim": 128,
        "expert_n_blocks": 0,
    }
    try:
        net.build_models(legacy_cfg)
        assert False, "expected ValueError"
    except ValueError as e:
        assert "expert_hidden_dim" in str(e)


def test_build_models_accepts_new_nested_shape_unchanged():
    """A dict that already has a 'stage1_model' key (the new shape) is
    passed through build_models without going through the legacy migration
    path at all."""
    cfg = {
        "pdg_vocab": PDG_VOCAB,
        "mat_vocab": MAT_VOCAB,
        "conditioning": {
            "out_dim": 32,
            "particle": {"type": "physical", "emb_dim": EMB_DIM, "n_layers": 1},
            "material": {"type": "physical", "emb_dim": EMB_DIM, "n_layers": 1},
        },
        "stage1_model": {
            "active": True,
            "generator": "flow",
            "hidden_dim": HIDDEN_DIM,
            "n_res_blocks": N_BLOCKS,
            "dropout": 0.0,
            "flow": {"time_dim": 16},
            "router": {"enabled": False},
        },
        "stage2_model": {
            "active": True,
            "decoder": "one_shot",
            "generator": "flow",
            "hidden_dim": HIDDEN_DIM,
            "n_res_blocks": N_BLOCKS,
            "dropout": 0.0,
            "k_max": K,
            "context_dim": 16,
            "n_sec": {"mode": "head"},
            "flow": {"time_dim": 16},
            "router": {"enabled": False, "tie_to_stage1": False},
        },
    }
    models = net.build_models(cfg)
    assert isinstance(models["stage1"], net.Stage1Model)
    assert isinstance(models["stage2"], net.Stage2OneShot)
    # Fresh v0.3.0 config, n_sec.owner defaults to "stage2": n_sec lives on stage 2.
    assert models["stage1"].n_sec_head is None
    assert models["stage2"].n_sec_head is not None