"""LoRA adapters on a frozen base backbone (``--lora-rank``), tiny random Qwen3.5 (CPU):

* zero-initialised adapters leave every head output bitwise unchanged;
* with trained (random) adapters, a field's outputs are bitwise the same alone, reordered
  and next to fields that run another adapter (per-row selector, batch-invariant);
* separate decide/extract adapters are isolated: changing one never moves the other
  capability's outputs;
* x-generate decodes adapter-off: the same text as the two-copy hybrid with the same base;
* training: only the selected adapter of each fork kind gets gradient, the base
  backbone none, and with ``context="base"`` the trunk carries no graph at all;
* save -> load round-trips the adapters (outputs bitwise equal).
"""

from __future__ import annotations

import copy

import pytest

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

from test_model import CONTEXT, SCHEMA, SNAPSHOT, _raw_outputs, _sub, tiny_lm  # noqa: E402
from test_train_mem import _examples  # noqa: E402

from deg.infer import fill  # noqa: E402
from deg.model import add_lora, load_model, load_tokenizer  # noqa: E402

FLAT = [n for n in SCHEMA["properties"] if n not in ("items", "reply")]
DECIDE = ["department", "urgent", "anger"]
EXTRACT = ["email", "order_id", "total", "charged_on"]


def _model(adapters=None, context="base", rank=4, randomize=True, lm=None):
    if SNAPSHOT is None:
        pytest.skip("Qwen3.5-0.8B tokenizer not cached")
    m = load_model(None, device="cpu", dtype=torch.float32, lm=lm or 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)
    if adapters is not None:
        m.backbone.requires_grad_(False)
        add_lora(m, rank, None, adapters, context)
        if randomize:
            g = torch.Generator().manual_seed(3)
            with torch.no_grad():
                for name, p in m.backbone.named_parameters():
                    if ".lora_B" in name:
                        p.copy_(torch.randn(p.shape, generator=g) * 0.05)
    return m


def _same(a, b, keys=None):
    keys = keys or a.keys()
    for k in keys:
        assert torch.equal(a[k], b[k]), k


def test_zero_adapters_are_the_base_model():
    plain = _model()
    zero = _model(["decide", "extract"], randomize=False)
    _same(_raw_outputs(plain, _sub(FLAT)), _raw_outputs(zero, _sub(FLAT)))


@pytest.mark.parametrize("adapters,context", [(["shared"], "base"),
                                              (["decide", "extract"], "base"),
                                              (["shared"], "shared")])
def test_adapter_rows_are_batch_invariant(adapters, context):
    m = _model(adapters, context)
    merged = _raw_outputs(m, _sub(FLAT))
    _same(merged, _raw_outputs(m, _sub(list(reversed(FLAT)))))
    for names in (DECIDE, EXTRACT, ["anger"], ["total"]):
        alone = _raw_outputs(m, _sub(names))
        _same(alone, merged, alone.keys())
    if adapters == ["decide", "extract"]:  # the adapters actually change the outputs
        plain = _raw_outputs(_model(), _sub(FLAT))
        assert any(not torch.equal(plain[k], merged[k]) for k in merged)


def test_separate_adapters_are_isolated():
    m = _model(["decide", "extract"])
    before = _raw_outputs(m, _sub(FLAT))
    r = m.lora["rank"]
    with torch.no_grad():  # perturb only the decide adapter's columns
        for name, p in m.backbone.named_parameters():
            if ".lora_B" in name:
                p[:, :r] += 0.1
    after = _raw_outputs(m, _sub(FLAT))
    _same(before, after, [f"/{n}" for n in EXTRACT])
    assert any(not torch.equal(before[f"/{n}"], after[f"/{n}"]) for n in DECIDE)


def test_generate_is_the_base_model():
    schema = _sub(["department", "order_id", "reply"])
    schema["properties"]["reply"] = {"type": "string", "x-generate": True, "maxLength": 40}
    lm = tiny_lm(0)
    base = copy.deepcopy(lm)
    for context in ("base", "shared"):
        m = _model(["shared"], context, lm=copy.deepcopy(lm))
        hyb = load_model(None, device="cpu", dtype=torch.float32, lm=copy.deepcopy(lm),
                         tokenizer=load_tokenizer(SNAPSHOT), generate="base",
                         base_lm=copy.deepcopy(base))
        got = fill(m, CONTEXT, schema)
        want_known = {k: got["value"][k] for k in ("department", "order_id")}
        # the hybrid decodes from the same known values with the same base weights
        hyb_plan = fill(hyb, CONTEXT, schema)
        if {k: hyb_plan["value"][k] for k in want_known} == want_known:
            assert got["value"]["reply"] == hyb_plan["value"]["reply"]
        if context == "base":  # one prefill: no second context read
            assert got["usage"]["prompt_tokens"] < hyb_plan["usage"]["prompt_tokens"]


def test_training_grads_follow_the_fork_kind():
    from deg.train import forward_losses

    m = _model(["decide", "extract"], "base")
    exs = _examples(m)
    m.backbone.train()
    m.heads.train()
    r = m.lora["rank"]
    try:
        hidden, state = m.trunk([exs[0].ctx])
        assert not hidden.requires_grad  # base context read: no trunk graph
        terms, _ = forward_losses(m, exs, 4096, 4)
        assert "lm" not in terms  # generate forks are not trained under LoRA
        terms["decide"].backward()
    finally:
        m.backbone.eval()
        m.heads.eval()
    got = {n: p.grad for n, p in m.backbone.named_parameters() if ".lora_B" in n}
    assert all(not p.requires_grad for n, p in m.backbone.named_parameters()
               if ".lora_" not in n)
    assert any(float(g[:, :r].abs().sum()) > 0 for g in got.values() if g is not None)
    assert all(g is None or float(g[:, r:].abs().sum()) == 0 for g in got.values())


def test_save_load_roundtrip(tmp_path):
    """Adapters (saved fp32) come back exactly; the backbone file is bf16, so outputs are
    compared between two loads (deterministic) and against zeroed adapters (applied)."""
    from deg.model import lora_state

    m = _model(["decide", "extract"], "base")
    m.save(tmp_path / "ck")
    assert (tmp_path / "ck" / "adapters.safetensors").exists()

    def load():
        return load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lm(5),
                          tokenizer=load_tokenizer(SNAPSHOT))

    back = load()
    assert back.lora == m.lora
    want = lora_state(m)
    got = lora_state(back)
    assert got.keys() == want.keys() and len(want) > 0
    for k in want:
        assert torch.equal(got[k], want[k].to(torch.bfloat16).float()) or \
            torch.equal(got[k], want[k]), k
    out = _raw_outputs(back, _sub(FLAT))
    _same(out, _raw_outputs(load(), _sub(FLAT)))
    with torch.no_grad():
        for name, p in back.backbone.named_parameters():
            if ".lora_B" in name:
                p.zero_()
    off = _raw_outputs(back, _sub(FLAT))
    assert any(not torch.equal(out[k], off[k]) for k in out)


def test_init_heads_loads_a_checkpoints_heads_onto_a_lora_model(tmp_path):
    """--init-heads: a full fine-tune's heads on the base backbone + fresh adapters."""
    from deg.train import init_heads

    ft = _model()  # randomized heads (the "full fine-tune")
    ft.save(tmp_path / "ft")
    m = _model(["shared"], "shared", randomize=False)
    with torch.no_grad():
        for p in m.heads.parameters():
            p.zero_()
    init_heads(m, tmp_path / "ft")
    for (k, a), b in zip(ft.heads.state_dict().items(), m.heads.state_dict().values()):
        assert torch.equal(a, b), k
    meta = tmp_path / "ft" / "deg.json"
    import json

    cfg = json.loads(meta.read_text())
    meta.write_text(json.dumps({**cfg, "backbone": "qwen3.5-2b"}))
    with pytest.raises(SystemExit):
        init_heads(m, tmp_path / "ft")