"""D/E/G task MoE on the Gemma 4 backbone (:func:`deg.model.add_task_moe`): FFN experts
per capability, deterministic routing by fork kind, shared attention / embeddings / PLE
(tiny random Gemma 4, CPU). Decide forks run D, extract / count / anchor forks and the
context run E (the layer's own FFN), x-generate (KNOWN VALUES, header, decoded tokens)
runs G; each expert's parameters get gradients only from its own losses."""

from __future__ import annotations

import json
import random

import pytest

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

from test_gemma_backbone import _gemma_model, _tok, tiny_gemma  # noqa: E402
from test_model import CONTEXT, SCHEMA, _raw_outputs, _sub  # noqa: E402

from deg.infer import fill  # noqa: E402
from deg.model import TaskRoute, add_task_moe, load_model, task_moe_params  # noqa: E402

DECIDE = ["department", "urgent", "anger"]
EXTRACT = ["email", "order_id", "total"]


@pytest.fixture(scope="module")
def pair():
    dense = _gemma_model(2)
    dense.meta["gen_layout"] = "chat"  # the task MoE writes in the chat turn (G route)
    moe = _gemma_model(2)
    add_task_moe(moe)
    return dense, moe


def _perturb(module, seed):
    g = torch.Generator().manual_seed(seed)
    with torch.no_grad():
        for p in module.parameters():
            p.add_(torch.randn(p.shape, generator=g) * 0.05)


def test_fresh_experts_are_the_dense_model(pair):
    dense, moe = pair
    want = _raw_outputs(dense, _sub(DECIDE + EXTRACT))
    got = _raw_outputs(moe, _sub(DECIDE + EXTRACT))
    for k in want:
        assert torch.equal(want[k], got[k]), k
    a, b = fill(dense, CONTEXT, SCHEMA), fill(moe, CONTEXT, SCHEMA)
    assert a["value"] == b["value"]
    assert moe.generator() == (moe, True, True)  # one prefill, G writes in the chat turn
    assert moe.gen_route(3) == TaskRoute("g")


def test_routing_isolates_the_experts():
    """Perturbing D changes only decide outputs; perturbing G only the generated text;
    perturbing E (the reader) changes everything that reads the context."""
    m = _gemma_model(2)
    add_task_moe(m)
    base = _raw_outputs(m, _sub(DECIDE + EXTRACT))
    reply = fill(m, CONTEXT, SCHEMA)["value"]["reply"]
    for layer in m.backbone.text.layers:
        _perturb(layer.mlp_d, 1)
    after_d = _raw_outputs(m, _sub(DECIDE + EXTRACT))
    assert all(torch.equal(after_d[f"/{n}"], base[f"/{n}"]) for n in EXTRACT)
    assert all(not torch.equal(after_d[f"/{n}"], base[f"/{n}"]) for n in DECIDE)
    m.graphs.clear()
    for layer in m.backbone.text.layers:
        _perturb(layer.mlp_g, 2)
    after_g = _raw_outputs(m, _sub(DECIDE + EXTRACT))
    for k in after_g:
        assert torch.equal(after_g[k], after_d[k]), k
    assert fill(m, CONTEXT, SCHEMA)["value"]["reply"] != reply
    for layer in m.backbone.text.layers:
        _perturb(layer.mlp, 3)
    after_e = _raw_outputs(m, _sub(DECIDE + EXTRACT))
    assert all(not torch.equal(after_e[k], after_g[k]) for k in after_e)


def test_each_expert_trains_only_on_its_own_losses():
    from deg.train import deg_example, forward_losses

    m = _gemma_model(2)
    m.meta["gen_layout"] = "chat"
    add_task_moe(m)
    a = CONTEXT.index("AB-1234")
    rec = {"uid": "t1", "schema": _sub(["department", "order_id", "reply"]),
           "context": CONTEXT,
           "value": {"department": "billing", "order_id": "AB-1234", "reply": "Refund issued."},
           "spans": {"/order_id": [a, a + 7]}, "span_alternatives": {}, "span_probs": {},
           "decide_probs": {}}
    ex = deg_example(m, rec, random.Random(0), shuffle_p=0.0)
    layers = m.backbone.text.layers

    def grads(keep):
        ex2 = type(ex)(**{**ex.__dict__, "items": [it for it in ex.items if it.kind in keep]})
        for p in m.backbone.parameters():
            p.grad = None
        m.backbone.train()
        try:
            losses, _ = forward_losses(m, [ex2], 4096, 16)
        finally:
            m.backbone.eval()
        sum(v for k, v in losses.items() if k in ("decide", "extract", "lm")).backward()
        return {x: any(p.grad is not None and bool(p.grad.abs().sum() > 0)
                       for lay in layers for p in getattr(lay, f"mlp{x}").parameters())
                for x in ("", "_d", "_g")}

    assert grads({"choice"}) == {"": True, "_d": True, "_g": False}  # E reads the context
    assert grads({"extract"}) == {"": True, "_d": False, "_g": False}
    assert grads({"generate"}) == {"": True, "_d": False, "_g": True}
    for p in m.backbone.parameters():
        p.grad = None


def test_task_moe_checkpoint_roundtrip_and_counts(tmp_path):
    m = _gemma_model(6)
    add_task_moe(m)
    with torch.no_grad():
        for p in m.backbone.lm.state_dict().values():
            if p.is_floating_point():
                p.copy_(p.bfloat16().float())
        for layer in m.backbone.text.layers:
            _perturb(layer.mlp_d, 4)
            for p in layer.mlp_d.parameters():
                p.copy_(p.bfloat16().float())
    m.save(tmp_path / "ck")
    meta = json.loads((tmp_path / "ck" / "deg.json").read_text())
    assert meta["moe"]["experts"] == ["d", "g"]
    other = tiny_gemma(3)
    with torch.no_grad():
        other.model.embed_tokens_per_layer.weight.copy_(
            m.backbone.text.embed_tokens_per_layer.weight)
    back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=other,
                      tokenizer=_tok())
    assert back.moe is not None and back.moe["experts"] == ["d", "g"]
    want = _raw_outputs(m, _sub(DECIDE + EXTRACT))
    got = _raw_outputs(back, _sub(DECIDE + EXTRACT))
    for k in want:
        assert torch.equal(want[k], got[k]), k
    c = task_moe_params(back)
    assert c["ffn_d"] == c["ffn_g"] == c["ffn_e"] > 0
    dense = sum(p.numel() for p in tiny_gemma(0).parameters())  # tied head counted once
    assert c["active_without_ple"] + c["ple"] == dense
    assert c["total"] == dense + c["ffn_d"] + c["ffn_g"]