"""deg-MoE (:func:`deg.model.add_moe`, :mod:`deg.moetrain`) on a tiny random Qwen3.5 (CPU):

* decide/extract outputs are bitwise the plain model's (W attached, random, never read);
* routing: R-only calls never read W (W poisoned with NaN changes nothing) and W calls
  never read R's FFN (R poisoned after the trunk changes nothing in generation);
* generate rows are batch-invariant; save/load round-trips; one train step moves only W.
"""

from __future__ import annotations

import pytest

torch = pytest.importorskip("torch")
pytest.importorskip("transformers")

from test_gentrain import RECORD  # noqa: E402
from test_hybrid_generate import SCHEMA  # noqa: E402
from test_model import CONTEXT, SNAPSHOT, _with_note, tiny_lm  # noqa: E402

from deg.gentrain import gen_items  # noqa: E402
from deg.infer import fill  # noqa: E402
from deg.model import add_moe, load_model, load_tokenizer, moe_params  # noqa: E402
from deg.moetrain import OffloadAdamW, batch_losses, micro_batches  # noqa: E402
from deg.request import FillOptions  # noqa: E402


def _plain():
    if SNAPSHOT is None:
        pytest.skip("Qwen3.5-0.8B tokenizer not cached")
    m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(0),
                   tokenizer=load_tokenizer(SNAPSHOT))
    torch.manual_seed(1)
    with torch.no_grad():
        for p in m.heads.parameters():
            p.normal_(0.0, 0.5)
    return m


def _moe():
    m = _plain()
    params = add_moe(m, tiny_lm(7), writer_rank=4)
    g = torch.Generator().manual_seed(5)
    with torch.no_grad():  # a non-trivial writer LoRA
        for p in params:
            if p.dim() == 2 and p.shape[1] == 4:
                p.copy_(torch.randn(p.shape, generator=g) * 0.2)
    return m, params


def _de(resp):
    return {k: (resp["value"][k], resp["meta"][f"/{k}"]) for k in ("department", "order_id")}


def test_decide_extract_bitwise_and_r_never_reads_w():
    plain = fill(_plain(), CONTEXT, SCHEMA)
    m, params = _moe()
    got = fill(m, CONTEXT, SCHEMA)
    assert _de(got) == _de(plain)
    with torch.no_grad():
        for layer in m.backbone.text.layers:
            for p in layer.mlp_w.parameters():
                p.fill_(float("nan"))
    poisoned = fill(m, CONTEXT, {**SCHEMA, "properties": {
        k: v for k, v in SCHEMA["properties"].items() if k != "reply"},
        "required": ["department", "order_id"]})
    assert _de(poisoned) == _de(plain)
    c = moe_params(m)
    assert c["w_ffn"] == c["r_ffn"] and c["active_writer"] < c["total"]


def test_w_never_reads_r_ffn_and_rows_are_batch_invariant():
    m, _ = _moe()
    ref = fill(m, CONTEXT, SCHEMA)
    assert isinstance(ref["value"]["reply"], str)
    # generation after the trunk: poison R's FFN only for W calls' benefit -> same text
    from deg import model as dm

    real = dm._expert

    def guarded(layer, sel, name="mlp"):
        e = real(layer, sel, name)
        if isinstance(sel, dm.WriteRoute):
            assert e is layer.mlp_w
        else:
            assert e is layer.mlp
        return e

    dm._expert = guarded
    try:
        assert fill(m, CONTEXT, SCHEMA) == ref
    finally:
        dm._expert = real
    for opts in (FillOptions(), FillOptions(temperature=0.9, seed=3)):
        alone = fill(m, CONTEXT, SCHEMA, opts)
        both = fill(m, CONTEXT, _with_note(SCHEMA), opts)
        assert both["value"]["reply"] == alone["value"]["reply"]
        assert both["meta"]["/reply"] == alone["meta"]["/reply"]


def test_save_load_roundtrip(tmp_path):
    m, _ = _moe()
    with torch.no_grad():  # the checkpoint stores the bodies in bf16
        for k, v in m.backbone.lm.state_dict().items():
            if ".lora_" not in k:
                v.copy_(v.to(torch.bfloat16).float())
    ref = fill(m, CONTEXT, SCHEMA)
    m.save(tmp_path / "ck")
    back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lm(3),
                      tokenizer=load_tokenizer(SNAPSHOT))
    assert back.moe == m.moe and back.generate_with == "moe"
    assert fill(back, CONTEXT, SCHEMA) == ref
    with pytest.raises(ValueError):
        load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lm(3),
                   tokenizer=load_tokenizer(SNAPSHOT), generate="base")


def test_one_train_step_moves_only_w():
    m, params = _moe()
    for p in params:
        p.requires_grad_(True)
    m.backbone.grad_ckpt = True
    frozen = {k: v.clone() for k, v in m.backbone.lm.state_dict().items()
              if ".mlp_w." not in k and ".lora_" not in k}
    w0 = [p.detach().clone() for p in params]
    items = gen_items(m, [RECORD], {"judge_drop": 0, "budget_drop": 0})
    (group,) = micro_batches(items, 4096, None)
    lora = [p.dim() == 2 and p.shape[1] == 4 or p.shape[0] == 4 for p in params]
    opt = OffloadAdamW(params, lora, 1e-3, 1e-3, torch.device("cpu"))
    nll, _, _, n = batch_losses(m, group)
    assert n == items[0].target_len and torch.isfinite(nll)
    (nll / n).backward()
    opt.step()
    for k, v in m.backbone.lm.state_dict().items():
        if k in frozen:
            assert torch.equal(v, frozen[k]), k
    assert any(not torch.equal(p, q) for p, q in zip(params, w0))