import copy

import pytest
import torch
from giant import config as gconfig
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
from giant.model.network import (
    AttentionHistory,
    ConditionEncoder,
    MarkovHistory,
    SinusoidalEmbedding,
    Stage1Model,
    Stage2Autoregressive,
    Stage2OneShot,
    build_models,
    cat_col_layout,
    stage2_trunk_sec_dim,
    stage2_type_dim,
)

PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}


def test_sinusoidal_embedding_shape():
    emb = SinusoidalEmbedding(64)
    t = torch.rand(16)
    assert emb(t).shape == (16, 64)


def test_sinusoidal_embedding_batch_1():
    emb = SinusoidalEmbedding(32)
    t = torch.tensor([0.5])
    assert emb(t).shape == (1, 32)


def test_stage1_model_output_shape():
    B = 8
    model = Stage1Model(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        n_sec_head_k_max=15,
    )
    x_t = torch.randn(B, 9)
    t = torch.rand(B)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [
            torch.randint(0, 5, (B,)),
            torch.randint(0, 3, (B,)),
        ],
        dim=1,
    )
    out = model(x_t, cond_cont, cond_cat, t=t)
    assert out.shape == (B, 9)


def test_stage1_model_gradients_flow():
    B = 4
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=PARTICLE_CFG,
        material_cfg=MATERIAL_CFG,
        hidden_dim=32,
        n_res_blocks=2,
        n_sec_head_k_max=15,
    )
    x_t = torch.randn(B, 9)
    t = torch.rand(B)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    # Both paths must be exercised to get gradients through all parameters.
    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_stage1_model_no_n_sec_head_by_default():
    """Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
    it moves to stage 2."""
    model = Stage1Model(
        pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
    )
    assert model.n_sec_head is None


# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------


def test_cat_col_layout_neither_onehot():
    assert cat_col_layout("physical", "embedding") == (None, None)


def test_cat_col_layout_particle_only():
    assert cat_col_layout("onehot", "physical") == (2, None)


def test_cat_col_layout_material_only():
    assert cat_col_layout("physical", "onehot") == (None, 2)


def test_cat_col_layout_both_onehot_particle_then_material():
    assert cat_col_layout("onehot", "onehot") == (2, 3)


def test_stage2_type_dim_physical_is_particle_phys_dim():
    assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM


def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
    assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
    assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16


def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
    k_max = 15
    assert (
        stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
        == k_max * SEC_SLOT_DIM
    )
    assert (
        stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
        == k_max * SEC_SLOT_DIM
    )


def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
    k_max = 15
    assert stage2_trunk_sec_dim(
        {"target": "onehot"}, "wgan", k_max, emb_dim=16
    ) == k_max * (CONT_SLOT_DIM + 16)


def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
    k_max = 15
    assert (
        stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
        == k_max * CONT_SLOT_DIM
    )


# --- ConditionEncoder onehot mode -------------------------------------------


def test_condition_encoder_onehot_forward_shape_and_gradients():
    B = 8
    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
    material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
    enc = ConditionEncoder(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=ONEHOT_PARTICLE_CFG,
        material_cfg=ONEHOT_MATERIAL_CFG,
        out_dim=32,
    )
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack(
        [
            torch.randint(0, 5, (B,)),
            torch.randint(0, 3, (B,)),
            torch.randint(0, particle_emb_dim, (B,)),
            torch.randint(0, material_emb_dim, (B,)),
        ],
        dim=1,
    )
    out = enc(cond_cont, cond_cat)
    assert out.shape == (B, 32)
    # onehot itself is unlearned, but the fusion MLP downstream still has
    # gradients — the encoder as a whole must still be trainable.
    out.sum().backward()
    assert enc.mlp[0].weight.grad is not None


def test_condition_encoder_onehot_is_a_true_one_hot_vector():
    """The onehot axis feeds a fixed, unlearned one-hot into the fusion MLP —
    verify the concatenated input segment really is one-hot, not e.g. an
    accidentally-learned embedding."""
    B = 4
    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
    enc = ConditionEncoder(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=ONEHOT_PARTICLE_CFG,
        material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
        out_dim=16,
    )
    cond_cont = torch.zeros(B, COND_DIM)
    idx = torch.tensor([0, 1, 2, 5])
    cond_cat = torch.stack(
        [
            torch.zeros(B, dtype=torch.long),
            torch.zeros(B, dtype=torch.long),
            idx.clamp(max=particle_emb_dim - 1),
        ],
        dim=1,
    )
    pdg_e = enc._particle_embed(cond_cont, cond_cat)
    assert pdg_e.shape == (B, particle_emb_dim)
    assert torch.all(pdg_e.sum(dim=-1) == 1.0)


# --- Stage2OneShot particle_type architecture --------------------------------


def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
    if target != "physical":
        particle_cfg = dict(particle_cfg)
        if target == "embedding":
            particle_cfg["type"] = "embedding"
    k_max = 5
    sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
    return Stage2OneShot(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=particle_cfg,
        material_cfg=MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        cond_out_dim=16,
        context_dim=8,
        sec_dim=sec_dim,
        generator=generator,
        k_max=k_max,
        particle_type_cfg={"target": target, "lambda": 1.0},
    )


def test_stage2_oneshot_physical_has_no_type_head_regardless_of_generator():
    assert _build_stage2("physical", "flow").type_head is None
    assert _build_stage2("physical", "wgan").type_head is None


def test_stage2_oneshot_onehot_flow_has_type_head():
    model = _build_stage2("onehot", "flow")
    assert model.type_head is not None


def test_stage2_oneshot_onehot_wgan_has_no_type_head():
    """Under wgan the type slice is folded into forward()'s own output and
    relaxed via ST-Gumbel by the trainer — no separate head needed."""
    model = _build_stage2("onehot", "wgan")
    assert model.type_head is None


def test_stage2_oneshot_embedding_flow_has_type_head():
    model = _build_stage2("embedding", "flow")
    assert model.type_head is not None


def test_stage2_oneshot_predict_type_shape():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    out = model.predict_type(cond_cont, cond_cat, stage1_out)
    assert out.shape == (B, k_max, emb_dim)


def test_stage2_oneshot_predict_type_raises_when_no_type_head():
    model = _build_stage2("physical", "flow")
    cond_cont = torch.randn(2, COND_DIM)
    cond_cat = torch.zeros(2, 2, dtype=torch.long)
    stage1_out = torch.randn(2, 9)
    try:
        model.predict_type(cond_cont, cond_cat, stage1_out)
        raise AssertionError("expected RuntimeError")
    except RuntimeError:
        pass


def test_stage2_oneshot_forward_shape_onehot_wgan():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    z = torch.randn(B, model.noise_dim)
    out = model(z, cond_cont, cond_cat, stage1_out)
    assert out.shape == (B, k_max * (CONT_SLOT_DIM + emb_dim))


def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    x_t = torch.randn(B, k_max * CONT_SLOT_DIM)
    t = torch.rand(B)
    out = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
    assert out.shape == (B, k_max * CONT_SLOT_DIM)


# --- MarkovHistory -----------------------------------------------------------


def test_markov_history_shape():
    hist = MarkovHistory(in_dim=7, out_dim=12)
    B, K = 3, 5
    feat = torch.randn(B, K, 7)
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    out = hist(feat, has_prev)
    assert out.shape == (B, K, 12)


def test_markov_history_uses_start_vector_when_no_prev():
    """Slot 0's own raw feature must be ignored — a learned start vector is
    substituted there instead (a reasonable default, see
    Stage2Autoregressive's docstring)."""
    hist = MarkovHistory(in_dim=4, out_dim=6)
    B, K = 2, 3
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    feat_a = torch.randn(B, K, 4)
    feat_b = feat_a.clone()
    feat_b[:, 0] = torch.randn(B, 4) * 100
    out_a = hist(feat_a, has_prev)
    out_b = hist(feat_b, has_prev)
    assert torch.allclose(out_a[:, 0], out_b[:, 0])
    assert torch.allclose(out_a[:, 1:], out_b[:, 1:])


# --- AttentionHistory (v0.3.0 step 7) ---------------------------------------


def test_attention_history_shape():
    hist = AttentionHistory(in_dim=7, out_dim=12, n_heads=2, n_layers=2)
    B, K = 3, 5
    feat = torch.randn(B, K, 7)
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    out = hist(feat, has_prev)
    assert out.shape == (B, K, 12)


def test_attention_history_uses_start_vector_when_no_prev():
    hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=1)
    B, K = 2, 3
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    feat_a = torch.randn(B, K, 4)
    feat_b = feat_a.clone()
    feat_b[:, 0] = torch.randn(B, 4) * 100
    out_a = hist(feat_a, has_prev)
    out_b = hist(feat_b, has_prev)
    assert torch.allclose(out_a[:, 0], out_b[:, 0], atol=1e-5)


def test_attention_history_is_causal():
    """Position i's output must not depend on feat at positions > i — unlike
    MarkovHistory (which only ever looks at position i itself, already
    trivially "causal"), this is AttentionHistory's actual contribution:
    seeing the full prefix 0..i-1, never anything later."""
    hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
    hist.eval()
    B, K = 2, 5
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    feat_a = torch.randn(B, K, 4)
    feat_b = feat_a.clone()
    feat_b[:, 3:] = torch.randn(B, K - 3, 4) * 100
    with torch.no_grad():
        out_a = hist(feat_a, has_prev)
        out_b = hist(feat_b, has_prev)
    assert torch.allclose(out_a[:, :3], out_b[:, :3], atol=1e-5)


def test_attention_history_step_matches_forward():
    """The incremental KV-cache path (`init_cache`/`step`,
    `giant/sample.py`'s AR loop) must reproduce `forward`'s parallel-pass
    output exactly, one position at a time."""
    hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
    hist.eval()
    B, K = 3, 6
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    feat = torch.randn(B, K, 4)
    with torch.no_grad():
        expected = hist(feat, has_prev)

        cache = hist.init_cache()
        outs = []
        for k in range(K):
            out_k, cache = hist.step(feat[:, k : k + 1], has_prev[:, k : k + 1], cache)
            outs.append(out_k)
        stepped = torch.cat(outs, dim=1)

    assert torch.allclose(stepped, expected, atol=1e-5)


# --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------


def _build_stage2_ar(
    target: str,
    generator: str,
    emb_dim: int = 6,
    k_max: int = 5,
    history: str = "markov",
) -> Stage2Autoregressive:
    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
    if target == "embedding":
        particle_cfg = dict(particle_cfg)
        particle_cfg["type"] = "embedding"
    return Stage2Autoregressive(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=particle_cfg,
        material_cfg=MATERIAL_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        cond_out_dim=16,
        context_dim=8,
        generator=generator,
        k_max=k_max,
        particle_type_cfg={"target": target, "lambda": 1.0},
        history=history,
    )


def _ar_inputs(B: int, K: int, hist_dim: int):
    history_feat = torch.randn(B, K, hist_dim)
    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    remaining_frac = torch.rand(B, K)
    slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1)
    return history_feat, has_prev, remaining_frac, slot_idx


def test_stage2_autoregressive_history_invalid_raises():
    with pytest.raises(ValueError):
        _build_stage2_ar("onehot", "wgan", history="bogus")


@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
@pytest.mark.parametrize("generator", ["wgan", "flow"])
@pytest.mark.parametrize("history", ["markov", "attention"])
def test_stage2_autoregressive_forward_shape(target, generator, history):
    B, K, emb_dim = 4, 5, 6
    model = _build_stage2_ar(
        target, generator, emb_dim=emb_dim, k_max=K, history=history
    )
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    type_dim = stage2_type_dim({"target": target}, emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
        B, K, CONT_SLOT_DIM + type_dim
    )
    token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
    if generator == "wgan":
        x_t = torch.randn(B, K, model.noise_dim)
        t = None
    else:
        x_t = torch.randn(B, K, token_dim)
        t = torch.rand(B, K)
    out = model(
        x_t,
        cond_cont,
        cond_cat,
        stage1_out,
        history_feat,
        has_prev,
        remaining_frac,
        slot_idx,
        t=t,
    )
    assert out.shape == (B, K, token_dim)


def test_stage2_autoregressive_predict_n_sec_shape():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=k_max)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    logits = model.predict_n_sec(cond_cont, cond_cat, stage1_out)
    assert logits.shape == (B, k_max + 1)


def test_stage2_autoregressive_predict_type_shape():
    B, K, emb_dim = 4, 5, 6
    model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
        B, K, CONT_SLOT_DIM + type_dim
    )
    out = model.predict_type(
        cond_cont,
        cond_cat,
        stage1_out,
        history_feat,
        has_prev,
        remaining_frac,
        slot_idx,
    )
    assert out.shape == (B, K, emb_dim)


@pytest.mark.parametrize("target,generator", [("physical", "flow"), ("onehot", "wgan")])
def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, generator):
    B, K, emb_dim = 2, 5, 6
    model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    type_dim = stage2_type_dim({"target": target}, emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
        B, K, CONT_SLOT_DIM + type_dim
    )
    with pytest.raises(RuntimeError):
        model.predict_type(
            cond_cont,
            cond_cat,
            stage1_out,
            history_feat,
            has_prev,
            remaining_frac,
            slot_idx,
        )


def test_stage2_autoregressive_gradients_flow_wgan_onehot():
    B, K, emb_dim = 4, 5, 6
    model = _build_stage2_ar("onehot", "wgan", emb_dim=emb_dim, k_max=K)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
        B, K, CONT_SLOT_DIM + type_dim
    )
    z = torch.randn(B, K, model.noise_dim)
    gen_out = model(
        z,
        cond_cont,
        cond_cat,
        stage1_out,
        history_feat,
        has_prev,
        remaining_frac,
        slot_idx,
    ).sum()
    nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
    (gen_out + nsec_out).backward()
    for name, p in model.named_parameters():
        assert p.grad is not None, f"no grad for {name}"


def test_stage2_autoregressive_gradients_flow_onehot():
    B, K, emb_dim = 4, 5, 6
    model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
        B, K, CONT_SLOT_DIM + type_dim
    )
    token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
    x_t = torch.randn(B, K, token_dim)
    t = torch.rand(B, K)
    flow_out = model(
        x_t,
        cond_cont,
        cond_cat,
        stage1_out,
        history_feat,
        has_prev,
        remaining_frac,
        slot_idx,
        t=t,
    ).sum()
    nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
    type_out = model.predict_type(
        cond_cont,
        cond_cat,
        stage1_out,
        history_feat,
        has_prev,
        remaining_frac,
        slot_idx,
    ).sum()
    (flow_out + nsec_out + type_out).backward()
    for name, p in model.named_parameters():
        assert p.grad is not None, f"no grad for {name}"


def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
    """`init_history_cache`/`history_step` (the incremental path
    `giant/sample.py`'s AR loop drives, one slot per call) must reproduce
    exactly what one parallel `self.history_encoder(history_feat, has_prev)`
    call over the whole shifted sequence would give at each position — the
    KV-cache correctness guarantee, exercised through `Stage2Autoregressive`
    itself rather than `AttentionHistory` in isolation
    (`test_attention_history_step_matches_forward` covers that lower layer)."""
    B, K, emb_dim = 3, 6, 6
    model = _build_stage2_ar(
        "physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
    )
    model.eval()
    type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
    hist_in_dim = CONT_SLOT_DIM + type_dim
    own_feat = torch.randn(B, K, hist_in_dim)  # token i's own raw feature
    has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
    history_feat = torch.cat(
        [torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
    )

    with torch.no_grad():
        expected = model.history_encoder(history_feat, has_prev_full)

        cache = model.init_history_cache()
        outs = []
        prev = torch.zeros(B, 1, hist_in_dim)
        for k in range(K):
            has_prev_k = torch.full((B, 1), k >= 1, dtype=torch.bool)
            hist_k, cache = model.history_step(prev, has_prev_k, cache)
            outs.append(hist_k)
            prev = own_feat[:, k : k + 1]
        stepped = torch.cat(outs, dim=1)

    assert torch.allclose(stepped, expected, atol=1e-5)


def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
    model = _build_stage2_ar("physical", "wgan", history="markov")
    assert model.init_history_cache() is None


# ── build_models: conditioning.share_stages ─────────────────────────────────


def _minimal_model_config(share_stages: bool) -> dict:
    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
    cfg["conditioning"]["share_stages"] = share_stages
    cfg["conditioning"]["particle"]["emb_dim"] = 4
    cfg["conditioning"]["material"]["emb_dim"] = 4
    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
    return {
        "pdg_vocab": 3,
        "mat_vocab": 2,
        "conditioning": cfg["conditioning"],
        "stage1_model": cfg["stage1_model"],
        "stage2_model": cfg["stage2_model"],
    }


def test_build_models_share_stages_true_shares_condition_encoder_instance():
    built = build_models(_minimal_model_config(share_stages=True))
    stage1, stage2 = built["stage1"], built["stage2"]
    assert stage1 is not None and stage2 is not None
    assert stage1.cond_enc is stage2.cond_enc


def test_build_models_share_stages_false_builds_independent_condition_encoders():
    built = build_models(_minimal_model_config(share_stages=False))
    stage1, stage2 = built["stage1"], built["stage2"]
    assert stage1 is not None and stage2 is not None
    assert stage1.cond_enc is not stage2.cond_enc


def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists():
    """The shared encoder's parameters must actually appear in both stages'
    own `.parameters()` — that's what makes each stage's independent
    optimizer include (and update) them, which is the actual mechanism behind
    "shared weights, forced common representation", not just object identity
    on `.cond_enc`."""
    built = build_models(_minimal_model_config(share_stages=True))
    stage1, stage2 = built["stage1"], built["stage2"]
    assert stage1 is not None and stage2 is not None

    shared_ids = {id(p) for p in stage1.cond_enc.parameters()}
    assert shared_ids
    assert shared_ids <= {id(p) for p in stage1.parameters()}
    assert shared_ids <= {id(p) for p in stage2.parameters()}