"""The LFM2.5 backbone (``--backbone lfm2.5-1.2b``; SPEC.md §12.8) on a tiny random LFM2
(CPU, exact torch reference kernels) with the real LFM2.5 tokenizer, plus GPU tests on the
real LFM2.5-1.2B-Instruct.

The model-contract tests of ``test_model`` / ``test_hybrid_generate`` / ``test_train_long``
are imported and collected here against this module's LFM fixtures (field independence
is bitwise alone / reordered / in bigger schemas / across requests, chunked prefill equals
the one-shot read, OOM splits are bitwise, hybrid generate decodes with the frozen base);
LFM-specific tests below check the functional stack against HF ``Lfm2ForCausalLM``,
forked fields against independent full forwards, the renderer's tokens, checkpoints and
training gradients.
"""

from __future__ import annotations

import random

import pytest

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

import test_hybrid_generate as thg  # noqa: E402
import test_model as tm  # noqa: E402
import test_train_long as ttl  # noqa: E402
from test_hybrid_generate import (  # noqa: E402,F401  (collected here with the LFM pair)
    test_hybrid_fill_many_matches_fill,
    test_hybrid_generate_rows_are_batch_invariant,
    test_hybrid_keeps_deg_fields_and_decodes_text_with_the_base,
)
from test_model import (  # noqa: E402,F401  (collected here with the LFM model)
    CONTEXT,
    _raw_outputs,
    _sub,
    test_blocked_span_logits_match_one_shot,
    test_consume_extends_state_without_holding_two_copies,
    test_field_outputs_are_bitwise_independent,
    test_fill_many_matches_fill_alone,
    test_fill_meta_is_identical_in_split_and_merged_schemas,
    test_fork_batch_split_on_oom_gives_the_same_outputs,
    test_generate_fields_are_independent_of_each_other,
    test_long_prefill_padded_batch_and_prefix_state,
    test_long_prefill_streams_in_chunks_without_quadratic_tensors,
    test_only_x_generate_fields_use_the_lm_head,
    test_per_field_generate_budget,
    test_sampled_generation_is_reproducible_with_seed,
    test_span_pointer_addresses_context_tokens_only,
)
from test_train_long import test_inference_long_read_unchanged_by_training_path  # noqa: E402,F401

from deg import compile_schema, validate_value  # noqa: E402
from deg.infer import fill  # noqa: E402
from deg.model import LFM2, base_snapshot, load_model, load_tokenizer  # noqa: E402

BACKBONE = "lfm2.5-1.2b"
try:
    LFM_SNAPSHOT = base_snapshot(BACKBONE)
except FileNotFoundError:  # pragma: no cover - needs the local HF cache
    LFM_SNAPSHOT = None


def tiny_lfm(seed: int = 0):
    from transformers import Lfm2Config, Lfm2ForCausalLM

    torch.manual_seed(seed)
    cfg = Lfm2Config(
        vocab_size=65536, hidden_size=64, intermediate_size=128,
        block_auto_adjust_ff_dim=False, num_hidden_layers=4,
        layer_types=["conv", "full_attention", "conv", "full_attention"],
        num_attention_heads=4, num_key_value_heads=2, conv_L_cache=3, conv_bias=False,
        norm_eps=1e-5, tie_embedding=True, max_position_embeddings=4096,
        rope_parameters={"rope_type": "default", "rope_theta": 10000.0},
        pad_token_id=0, bos_token_id=1, eos_token_id=7)
    lm = Lfm2ForCausalLM(cfg).float().eval()
    with torch.no_grad():  # random norms so the stack is not near-identity
        for name, p in lm.named_parameters():
            if "norm" in name:
                p.normal_(1.0, 0.3)
    return lm


def _tok():
    if LFM_SNAPSHOT is None:
        pytest.skip("LFM2.5-1.2B-Instruct tokenizer not cached")
    return load_tokenizer(backbone=BACKBONE)


def _lfm_model(seed: int = 0, **kw):
    m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lfm(seed),
                   tokenizer=_tok(), **kw)
    torch.manual_seed(1)
    with torch.no_grad():
        for p in m.heads.parameters():
            p.normal_(0.0, 0.5)
    return m


@pytest.fixture(scope="module")
def model():
    return _lfm_model()


@pytest.fixture(scope="module")
def pair():
    """(deg-only, hybrid) LFM models: the hybrid's frozen base is another random LFM
    with the same (frozen) embeddings, as deg training leaves them."""
    def make(generate):
        base, deg = tiny_lfm(7), tiny_lfm(0)
        with torch.no_grad():
            deg.model.embed_tokens.weight.copy_(base.model.embed_tokens.weight)
        m = load_model(None, device="cpu", dtype=torch.float32, lm=deg, tokenizer=_tok(),
                       generate=generate, base_lm=base)
        torch.manual_seed(1)
        with torch.no_grad():
            for p in m.heads.parameters():
                p.normal_(0.0, 0.5)
        return m

    return make("self"), make("base")


# ----------------------------------------------------------------------- identity


def test_lfm_renderer_and_meta(model):
    rend = model.renderer
    assert model.backbone.fam is LFM2 and model.meta["backbone"] == BACKBONE
    assert model.backbone.types == ["rec", "attn", "rec", "attn"]
    ctx = rend.context(CONTEXT)
    tok = rend.tokenizer
    assert ctx.ids[0] == tok.convert_tokens_to_ids("<|startoftext|>")  # BOS
    assert tok.convert_ids_to_tokens(ctx.ids[1]) == "<|im_start|>"
    assert rend.pad_id == tok.convert_tokens_to_ids("<|pad|>")
    assert rend.eos_id == tok.convert_tokens_to_ids("<|im_end|>")
    assert rend.gen_turn == "<|im_end|>\n<|im_start|>assistant\n"  # no <think> block
    # every context token maps back to its text
    for i, (a, b) in enumerate(ctx.offsets[:20]):
        assert tok.decode([ctx.ids[ctx.ctx_start + i]]).strip() == CONTEXT[a:b].strip()


def test_lfm_backbone_matches_hf_and_fork_equals_concatenation(model):
    """The functional stack equals HF Lfm2Model on the concatenated read; continuing a
    state (dense and padded rows) equals reading the whole sequence."""
    rend = model.renderer
    a = list(rend.context(CONTEXT).ids)
    b = rend.enc("FIELD /order_id\nKIND: copy\nANSWER:")
    with torch.no_grad():
        ref = model.backbone.text(torch.tensor([a + b])).last_hidden_state[0]
        h1, st = model.backbone(torch.tensor([a]), torch.tensor([len(a)]))
        h2, _ = model.backbone(torch.tensor([b]), torch.tensor([len(b)]), st)
        pad = torch.full((2, len(a) + 5), rend.pad_id)
        pad[0, :len(a)] = torch.tensor(a)
        pad[1, :3] = torch.tensor(a[:3])
        _, st2 = model.backbone(pad, torch.tensor([len(a), 3]))
        fork = torch.full((1, len(b) + 9), rend.pad_id)
        fork[0, :len(b)] = torch.tensor(b)
        h3, _ = model.backbone(fork, torch.tensor([len(b)]), st2.select(torch.tensor([0])))
    torch.testing.assert_close(torch.cat([h1[0], h2[0]]), ref, atol=2e-5, rtol=2e-5)
    torch.testing.assert_close(h3[0, :len(b)], h2[0], atol=2e-5, rtol=2e-5)
    # conv layers carry no recurrent state, only a 2-token tail of B*x
    assert all(r.shape == (1, 0) and t.shape == (1, 2, 64) for r, t in st.lin)


def test_lfm_forked_fields_equal_independent_full_forward(model):
    """Each fork row (prefix read by reference, run in one batch with other fields) equals
    an independent HF full forward over context + that field's query; and a row's fork
    hidden states are bitwise the same alone and in the batch."""
    rend = model.renderer
    ctx = list(rend.context(CONTEXT).ids)
    queries = [rend.enc(t) for t in (
        "FIELD /order_id\nKIND: copy\nANSWER:",
        "FIELD /department\nOPTIONS:\n- billing\n- technical\nANSWER:",
        "ARRAY /items\nANSWER:",
        "FIELD /total\nTITLE: total charged\nKIND: number\nCONSTRAINTS: {}\nANSWER:")]
    bb = model.backbone

    def forks(rows):
        width = max(len(q) for q in rows)
        ids = torch.full((len(rows), width), rend.pad_id)
        for i, q in enumerate(rows):
            ids[i, :len(q)] = torch.tensor(q)
        _, st = bb(torch.tensor([ctx]), torch.tensor([len(ctx)]))
        fs = bb.fork(st, len(rows), width)
        return bb.fork_forward(ids, torch.tensor([len(q) for q in rows]), fs)

    with torch.no_grad():
        batch = forks(queries)
        for i, q in enumerate(queries):
            ref = bb.text(torch.tensor([ctx + q])).last_hidden_state[0, len(ctx):]
            torch.testing.assert_close(batch[i, :len(q)], ref, atol=2e-5, rtol=2e-5)
            alone = forks([q])
            assert torch.equal(alone[0, :len(q)], batch[i, :len(q)]), i


def test_lfm_fork_decode_continues_state(model):
    """keep=True fork passes (header, then one token at a time: the generate decode path)
    equal the full read of context + header + tokens."""
    rend = model.renderer
    bb = model.backbone
    ctx = list(rend.context(CONTEXT).ids)
    head = rend.enc("WRITE /reply\nVALUE:\n")
    more = rend.enc(" Refund issued today")
    with torch.no_grad():
        _, st = bb(torch.tensor([ctx]), torch.tensor([len(ctx)]))
        fs = bb.fork(st, 2, len(head) + len(more) + 1)
        ids = torch.tensor([head, head])
        bb.fork_forward(ids, torch.tensor([len(head)] * 2), fs, keep=True)
        outs = []
        for t in more:
            outs.append(bb.fork_forward(torch.tensor([[t], [t]]), torch.tensor([1, 1]), fs,
                                        keep=True)[:, 0])
        ref = bb.text(torch.tensor([ctx + head + more])).last_hidden_state[0, -len(more):]
    got = torch.stack(outs, 1)
    torch.testing.assert_close(got[0], ref, atol=2e-5, rtol=2e-5)
    assert torch.equal(got[0], got[1])


def test_lfm_chunked_checkpointed_trunk_equals_dense(model, monkeypatch):
    """Training trunks beyond DENSE_PREFILL_MAX (checkpointed chunks) equal the dense
    checkpointed read in hidden states, state, forks and gradients (test_train_long)."""
    emb = model.backbone.text.embed_tokens.weight
    model.backbone.grad_ckpt = True
    emb.requires_grad_(False)
    try:
        ttl.test_chunked_checkpointed_trunk_equals_dense(model, monkeypatch)
    finally:
        model.backbone.grad_ckpt = False
        emb.requires_grad_(True)
        for p in model.backbone.parameters():
            p.grad = None


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

    a = CONTEXT.index("AB-1234")
    rec = {"uid": "t1", "schema": _sub(["department", "order_id", "items", "reply"]),
           "context": CONTEXT,
           "value": {"department": "billing", "order_id": "AB-1234",
                     "items": [{"name": "chairs", "qty": 2}], "reply": "Refund issued."},
           "spans": {"/order_id": [a, a + 7],
                     "/items/0/name": [CONTEXT.index("chairs"),
                                       CONTEXT.index("chairs") + 6]},
           "span_alternatives": {}, "span_probs": {},
           "decide_probs": {"/department": [0.8, 0.15, 0.05]}}
    ex = deg_example(model, rec, random.Random(0), shuffle_p=0.0)
    assert sorted(it.kind for it in ex.items) == [
        "anchor", "choice", "extract", "extract", "generate", "score"]
    model.backbone.train()
    try:
        losses, _ = forward_losses(model, [ex], 4096, 16)
    finally:
        model.backbone.eval()
    total = sum(losses[k] for k in ("decide", "extract", "anchor", "lm"))
    assert torch.isfinite(total)
    total.backward()
    layers = model.backbone.text.layers
    assert layers[0].conv.in_proj.weight.grad is not None
    assert layers[1].self_attn.q_proj.weight.grad is not None
    assert layers[1].feed_forward.w1.weight.grad is not None
    for p in [*model.heads.parameters(), *model.backbone.parameters()]:
        p.grad = None


def test_lfm_checkpoint_roundtrip_keeps_backbone(tmp_path):
    model = _lfm_model(5)
    with torch.no_grad():  # checkpoints store the backbone in bf16
        for p in model.backbone.parameters():
            p.copy_(p.bfloat16().float())
    model.save(tmp_path / "ck")
    meta = __import__("json").loads((tmp_path / "ck" / "deg.json").read_text())
    assert meta["backbone"] == BACKBONE
    back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lfm(3),
                      tokenizer=_tok())
    assert back.backbone.fam is LFM2
    want = _raw_outputs(model, _sub(["order_id", "department"]))
    got = _raw_outputs(back, _sub(["order_id", "department"]))
    for k in want:
        assert torch.equal(want[k], got[k]), k
    with pytest.raises(ValueError):
        load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lfm(3),
                   tokenizer=_tok(), backbone="qwen3.5-0.8b")
    with pytest.raises(ValueError):  # a Qwen base cannot be the LFM model's generator
        from deg.model import attach_base

        attach_base(back, tm.tiny_lm(0))


def test_lfm_hybrid_base_shares_embeddings(pair):
    _, hyb = pair
    assert hyb.base is not None and hyb.base.shared_embeddings and hyb.base.fam is LFM2
    emb = hyb.backbone.text.embed_tokens.weight
    assert hyb.base.text.embed_tokens.weight is emb and hyb.base.lm.lm_head.weight is emb
    resp = fill(hyb, CONTEXT, thg.SCHEMA)
    validate_value(compile_schema(thg.SCHEMA), resp["value"])


# ---------------------------------------------------------------------------- GPU


def _gpu_free_gib() -> float:
    if not torch.cuda.is_available():
        return 0.0
    return torch.cuda.mem_get_info()[0] / 2**30


@pytest.mark.skipif(LFM_SNAPSHOT is None or _gpu_free_gib() < 4.0,
                    reason="needs CUDA with ~4 GiB free and the LFM2.5-1.2B snapshot")
def test_gpu_lfm_real_model_bitwise_independence_and_hf_match():
    """Real LFM2.5-1.2B in bf16 on CUDA: the fork path (bi_linear GEMMs, fork_attention,
    Triton LFM RMSNorm, elementwise short conv) gives bitwise identical raw outputs alone,
    reordered, inside a bigger schema and across requests (fill_many); the hybrid's base
    text is identical alone and next to another generate field; the functional trunk
    matches HF within bf16 noise."""
    m = load_model(None, device="cuda", backbone=BACKBONE, generate="base")
    rend = m.renderer
    ids = torch.tensor([rend.context(CONTEXT).ids], device="cuda")
    with torch.no_grad():
        ref = m.backbone.text(ids).last_hidden_state[0].float()
        mine, _ = m.backbone(ids, torch.tensor([ids.shape[1]], device="cuda"))
    cos = torch.nn.functional.cosine_similarity(mine[0].float(), ref, -1)
    assert float(cos.min()) > 0.99, float(cos.min())
    flat = [n for n in tm.SCHEMA["properties"] if n not in ("items", "reply")]
    merged = _raw_outputs(m, _sub(flat))
    reordered = _raw_outputs(m, _sub(list(reversed(flat))))
    big = {"type": "object", "properties": {
        **{f"f{i}": {"type": ["string", "null"], "description": f"filler {i}"}
           for i in range(12)}, **_sub(flat)["properties"]},
        "required": _sub(flat)["required"]}
    big_out = _raw_outputs(m, big)
    for name in flat:
        alone = _raw_outputs(m, _sub([name]))[f"/{name}"]
        for other in (merged, reordered, big_out):
            assert torch.equal(alone, other[f"/{name}"]), name
    resp = fill(m, CONTEXT, tm.SCHEMA)
    validate_value(compile_schema(tm.SCHEMA), resp["value"])
    both = fill(m, CONTEXT, tm._with_note())
    assert both["value"]["reply"] == resp["value"]["reply"]
    assert both["meta"]["/reply"] == resp["meta"]["/reply"]
    from deg.infer import fill_many

    other = "Dana again: the desk arrived broken. Order AB-9 total $80."
    jobs = [(CONTEXT, tm.SCHEMA, None), (other, tm._with_note(), None),
            (other, _sub(["department", "items"]), None)]
    for job, got in zip(jobs, fill_many(m, jobs)):
        assert got == fill(m, *job)


@pytest.mark.skipif(LFM_SNAPSHOT is None or _gpu_free_gib() < 4.0,
                    reason="needs CUDA with ~4 GiB free and the LFM2.5-1.2B snapshot")
def test_gpu_lfm_chunked_prefill_matches(monkeypatch):
    """Real LFM2.5-1.2B on CUDA: the chunked prefill (flash-style fork kernel over the K/V
    buffer, head_dim 64, 4 query heads per K/V head) matches the one-shot read within bf16
    noise for padded rows."""
    import deg.model as dm

    m = load_model(None, device="cuda", backbone=BACKBONE)
    rend = m.renderer
    a = rend.context(" ".join([CONTEXT] * 40)).ids
    b = rend.context(" ".join([CONTEXT] * 15)).ids
    x = torch.full((2, len(a)), rend.pad_id, device="cuda")
    x[0], x[1, :len(b)] = torch.tensor(a), torch.tensor(b)
    n = torch.tensor([len(a), len(b)], device="cuda")
    with torch.no_grad():
        h_one, s_one = m.backbone(x, n)
        monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 512)
        monkeypatch.setattr(dm, "PREFILL_CHUNK", 512)
        h, s = m.backbone(x, n)
    for row, L in enumerate((len(a), len(b))):
        cos = torch.nn.functional.cosine_similarity(h[row, :L].float(),
                                                    h_one[row, :L].float(), -1)
        assert float(cos.min()) > 0.99 and float(cos.mean()) > 0.9995, row
    assert torch.equal(s.kv_mask, s_one.kv_mask)