"""Model/runtime contract tests on a tiny random Qwen3.5 (CPU, exact torch reference
kernels) plus one GPU test on the real base model.

* the functional backbone equals the HF forward, and a fork equals the concatenated read;
* field independence is bitwise: a field's head outputs are identical alone, reordered,
  next to other fields, and in split vs merged schemas;
* the span pointer only addresses context tokens;
* only x-generate fields reach the LM head;
* responses validate against the schema; constrained decoding respects length/pattern.
"""

from __future__ import annotations

import random

import pytest

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

from deg import compile_schema, validate_value  # noqa: E402
from deg.infer import Constraint, field_query, fill  # noqa: E402
from deg.model import (  # noqa: E402
    base_snapshot,
    load_model,
    load_tokenizer,
    token_span,
    trim_span,
)
from deg.request import FillOptions  # noqa: E402
from deg.response import FillError  # noqa: E402

try:
    SNAPSHOT = base_snapshot()
except FileNotFoundError:  # pragma: no cover - needs the local HF cache
    SNAPSHOT = None

needs_tokenizer = pytest.mark.skipif(SNAPSHOT is None, reason="Qwen3.5-0.8B not cached")

CONTEXT = ("Ticket from dana@example.com: order AB-1234 was charged twice on 2026-03-05, "
           "total $1,250.50. Please refund today, I am furious. Items: 2 chairs, 1 desk.")
SCHEMA = {
    "type": "object",
    "properties": {
        "department": {"enum": ["billing", "technical", "sales"],
                       "description": "Which team should handle this?"},
        "urgent": {"type": "boolean"},
        "anger": {"type": "integer", "minimum": 0, "maximum": 4},
        "email": {"type": ["string", "null"], "format": "email"},
        "order_id": {"type": ["string", "null"]},
        "total": {"type": ["number", "null"]},
        "charged_on": {"type": ["string", "null"], "format": "date"},
        "items": {"type": "array", "maxItems": 4,
                  "items": {"type": "object",
                            "properties": {"name": {"type": ["string", "null"]},
                                           "qty": {"type": "integer", "minimum": 1,
                                                   "maximum": 5}},
                            "required": ["name", "qty"]}},
        "reply": {"type": "string", "x-generate": True, "maxLength": 40, "minLength": 3,
                  "pattern": "^[^{}]*$"},
    },
    "required": ["department", "urgent", "anger", "email", "order_id", "total", "charged_on",
                 "items", "reply"],
}


def tiny_lm(seed: int = 0):
    from transformers import Qwen3_5ForCausalLM, Qwen3_5TextConfig

    torch.manual_seed(seed)
    cfg = Qwen3_5TextConfig(
        vocab_size=248320, hidden_size=64, intermediate_size=128, num_hidden_layers=4,
        layer_types=["linear_attention", "linear_attention", "linear_attention",
                     "full_attention"],
        num_attention_heads=2, num_key_value_heads=1, head_dim=32,
        linear_num_key_heads=2, linear_num_value_heads=2, linear_key_head_dim=16,
        linear_value_head_dim=16, linear_conv_kernel_dim=4, tie_word_embeddings=True,
        rope_parameters={"rope_type": "default", "rope_theta": 10000.0,
                         "partial_rotary_factor": 0.25, "mrope_section": [2, 1, 1],
                         "mrope_interleaved": True},
    )
    lm = Qwen3_5ForCausalLM(cfg).float().eval()
    with torch.no_grad():  # random norms/gates so the trunk is not near-identity
        for name, p in lm.named_parameters():
            if "norm" in name or "A_log" in name or "dt_bias" in name:
                p.normal_(0.0, 0.3)
    return lm


@pytest.fixture(scope="module")
def model():
    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(),
                   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


# ----------------------------------------------------------------------- pure helpers


def test_trim_span_strips_token_punctuation():
    text = 'Buy "Desk (White)." now, [Marcus Thorne] and Inc.,'
    a = text.index('"')
    assert text.__getitem__(slice(*trim_span(text, a, text.index(" now")))) == "Desk (White)."
    a = text.index("[")
    assert text.__getitem__(slice(*trim_span(text, a, text.index(" and")))) == "Marcus Thorne"
    a = text.index("Inc")
    assert text.__getitem__(slice(*trim_span(text, a, len(text)))) == "Inc."


def _cut(text: str, sub: str, extra: int = 0) -> str:
    a = text.index(sub)
    return text[slice(*trim_span(text, a, a + len(sub) + extra))]


def test_trim_span_keeps_value_stars_and_strips_wrappers():
    # masked card number: the '*' run is the value, not markdown (run-1 s1-0002479)
    t = "on card ****4821. It"
    assert _cut(t, "****4821") == "****4821"
    # markdown emphasis is a symmetric wrapper, also when the token span stops short
    t = "Status: **bold** today"
    assert _cut(t, "**bold**") == "bold"
    assert _cut(t, "**bold") == "bold"
    assert _cut(t, "bold**") == "bold"
    assert _cut("x *it*, y", "*it*,") == "it"
    # quotes are stripped as a pair
    assert _cut('vendor "ACME Corp" signed', '"ACME Corp"') == "ACME Corp"
    # a trailing '*' that belongs to the value (one-sided) is kept
    assert _cut("Grade: AB*, final", "AB*,") == "AB*"


def _pointer_logits(offsets, i, j):
    T = len(offsets)
    start = torch.full((T + 1,), -20.0)
    end = torch.full((T + 1,), -20.0)
    start[i] = end[j] = 10.0
    return start, end


def test_extract_result_tries_raw_span_when_trimmed_fails_constraints():
    from deg.infer import extract_result
    from deg.model import EncodedContext

    def run(ctx_text, prop, value_sub):
        fld = compile_schema({"type": "object", "properties": {"v": prop},
                              "required": ["v"]}).fields[0]
        words = [(m.start(), m.end()) for m in __import__("re").finditer(r"\S+", ctx_text)]
        a = ctx_text.index(value_sub)
        i, j = token_span(tuple(words), a, a + len(value_sub))
        ctx = EncodedContext(ids=tuple(range(len(words))), ctx_start=0, offsets=tuple(words))
        start, end = _pointer_logits(words, i, j)
        return extract_result(fld, ctx_text, ctx, start, end, torch.tensor(1.0),
                              FillOptions()).value

    card = {"type": "string", "pattern": "^\\*{4}[0-9]{4}$", "minLength": 8}
    assert run("Paid on card ****4821, it cleared.", card, "****4821,") == "****4821"
    # quotes that are part of the value under a pattern: trimmed fails, raw passes
    quoted = {"type": "string", "pattern": '^"[A-Z]+"$'}
    assert run('Use code "ABC" now', quoted, '"ABC"') == '"ABC"'
    # unconstrained: trimmed wins
    assert run('vendor "ACME Corp" signed', {"type": "string"}, '"ACME Corp"') == "ACME Corp"
    assert run("Grade: AB*, final", {"type": "string", "pattern": "^[A-Z]+\\*$"},
               "AB*,") == "AB*"


def test_constraint_prefix_and_final():
    fld = compile_schema({"type": "object", "properties": {
        "code": {"type": "string", "x-generate": True, "pattern": "^[A-Z]{2}-\\d+$",
                 "maxLength": 8, "minLength": 4}}}).fields[0]
    rule = Constraint(fld)
    assert rule.prefix_ok("A") and rule.prefix_ok("AB-") and rule.prefix_ok("AB-12")
    assert not rule.prefix_ok("a") and not rule.prefix_ok("AB-123456")
    assert not rule.final_ok("AB-") and rule.final_ok("AB-12")


@needs_tokenizer
def test_token_span_covers_char_span():
    tok = load_tokenizer(SNAPSHOT)
    enc = tok(CONTEXT, add_special_tokens=False, return_offsets_mapping=True)
    offsets = tuple(tuple(x) for x in enc["offset_mapping"])
    a = CONTEXT.index("AB-1234")
    i, j = token_span(offsets, a, a + 7)
    assert offsets[i][0] <= a and offsets[j][1] >= a + 7
    got = trim_span(CONTEXT, offsets[i][0], offsets[j][1])
    assert CONTEXT[got[0]:got[1]] == "AB-1234"


# ------------------------------------------------------------------------- backbone


def test_backbone_matches_hf_and_fork_equals_concatenation(model, monkeypatch):
    from transformers.models.qwen3_5 import modeling_qwen3_5 as hf

    from deg.model import delta_rule_reference

    def chunk_ref(q, k, v, g, beta, initial_state=None, output_final_state=False, **_):
        return delta_rule_reference(q, k, v, g, beta, initial_state)

    monkeypatch.setattr(hf, "torch_chunk_gated_delta_rule", chunk_ref)  # HF CPU path
    rend = model.renderer
    a = rend.enc(CONTEXT)
    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)
        # padded rows (trunk and fork) give the same fork result
        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])))
    mine = torch.cat([h1[0], h2[0]])
    torch.testing.assert_close(mine, ref, atol=2e-4, rtol=2e-4)
    torch.testing.assert_close(h3[0, :len(b)], h2[0], atol=2e-4, rtol=2e-4)


# --------------------------------------------------------------------- independence


def _raw_outputs(model, schema):
    """pointer -> raw head tensors (logits or start/end) for every parallel leaf."""
    from deg.infer import LeafInstance, walk

    plan = compile_schema(schema)
    ectx = model.renderer.context(CONTEXT)
    with torch.no_grad():
        hidden, state = model.trunk([ectx])
        leaves = [x for x in walk(plan.root, {}) if isinstance(x, LeafInstance)
                  and x.field.phase == "parallel"]
        queries = [model.renderer.query(field_query(x.field, x.ptr, list(x.items)))
                   for x in leaves]
        outs = model.answer_batched(state, hidden, ectx, queries)
    return {x.ptr: (o.logits if o.logits is not None else torch.cat([o.start, o.end]))
            for x, o in zip(leaves, outs)}


def _sub(names, schema=SCHEMA):
    return {"type": "object",
            "properties": {n: schema["properties"][n] for n in names},
            "required": [n for n in names if n in schema["required"]]}


def test_field_outputs_are_bitwise_independent(model):
    flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")]
    merged = _raw_outputs(model, _sub(flat))
    assert len(merged) == len(flat)
    reordered = _raw_outputs(model, _sub(list(reversed(flat))))
    # more fields than one slot batch, so neighbours change across batches too
    padded = {n: {"type": ["string", "null"], "description": f"filler {n}"}
              for n in ("f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "f9")}
    big_schema = {"type": "object", "properties": {**padded, **_sub(flat)["properties"]},
                  "required": _sub(flat)["required"]}
    big = _raw_outputs(model, big_schema)
    assert len(big) > 8
    for name in flat:
        alone = _raw_outputs(model, _sub([name]))
        ptr = f"/{name}"
        for other in (merged, reordered, big):
            assert torch.equal(alone[ptr], other[ptr]), name


def test_fill_meta_is_identical_in_split_and_merged_schemas(model):
    names = ["department", "urgent", "anger", "email", "order_id", "total", "charged_on"]
    merged = fill(model, CONTEXT, _sub(names))
    for half in (names[:3], names[3:]):
        part = fill(model, CONTEXT, _sub(half))
        for n in half:
            assert part["meta"][f"/{n}"] == merged["meta"][f"/{n}"]
            assert part["value"].get(n, "<omitted>") == merged["value"].get(n, "<omitted>")


# ------------------------------------------------------------------------ pointer


def test_span_pointer_addresses_context_tokens_only(model):
    ectx = model.renderer.context(CONTEXT)
    fld = compile_schema(_sub(["order_id"])).fields[0]
    q = model.renderer.query(field_query(fld, "/order_id", []))
    with torch.no_grad():
        hidden, state = model.trunk([ectx])
        (out,) = model.answer_batched(state, hidden, ectx, [q])
    # one logit per context token + the null outcome; prefix/suffix/query tokens excluded
    assert out.start.shape[0] == len(ectx.offsets) + 1 == out.end.shape[0]
    assert len(ectx.offsets) < len(ectx.ids)
    for _ in range(5):
        resp = fill(model, CONTEXT, _sub(["order_id", "email", "total", "charged_on"]))
        for ptr, meta in resp["meta"].items():
            if meta.get("span") is not None:
                a, b = meta["span"]
                assert 0 <= a < b <= len(CONTEXT)


# ----------------------------------------------------------------------- generation


def test_only_x_generate_fields_use_the_lm_head(model, monkeypatch):
    calls = []
    real = model.backbone.lm_logits

    def spy(hidden):
        calls.append(hidden.shape)
        return real(hidden)

    monkeypatch.setattr(model.backbone, "lm_logits", spy)
    no_gen = {**SCHEMA, "properties": {k: v for k, v in SCHEMA["properties"].items()
                                       if k != "reply"},
              "required": [r for r in SCHEMA["required"] if r != "reply"]}
    resp = fill(model, CONTEXT, no_gen)
    assert calls == [] and resp["usage"]["generated_tokens"] == 0
    resp = fill(model, CONTEXT, SCHEMA)
    gen = [p for p, m in resp["meta"].items() if m.get("mode") == "generate"]
    assert gen == ["/reply"]
    assert len(calls) == resp["meta"]["/reply"]["generated_tokens"] > 0
    assert resp["usage"]["generated_tokens"] == len(calls)


def test_response_validates_and_generation_respects_constraints(model):
    for seed in range(3):
        opts = FillOptions(temperature=0.8, seed=seed) if seed else FillOptions()
        resp = fill(model, CONTEXT, SCHEMA, opts)
        validate_value(compile_schema(SCHEMA), resp["value"])
        reply = resp["value"]["reply"]
        assert 3 <= len(reply) <= 40 and "{" not in reply and "}" not in reply
        assert resp["model"] == "deg-qwen3.5-0.8b-0.1"
        assert set(resp["value"]) >= {"department", "urgent", "anger", "items", "reply"}


def test_sampled_generation_is_reproducible_with_seed(model):
    opts = FillOptions(temperature=1.0, seed=7)
    a = fill(model, CONTEXT, SCHEMA, opts)
    b = fill(model, CONTEXT, SCHEMA, opts)
    assert a == b


def test_required_field_without_any_parsable_span_is_a_fill_error(model):
    schema = {"type": "object", "properties": {"when": {"type": "string",
                                                        "format": "date-time"}},
              "required": ["when"]}
    with pytest.raises(FillError):
        fill(model, "no timestamps in here at all", schema)


def _with_note(schema=SCHEMA, budget=None):
    note = {"type": "string", "x-generate": True, "maxLength": 30,
            "description": "internal note for the agent"}
    if budget is not None:
        note["x-max-tokens"] = budget
    return {**schema, "properties": {**schema["properties"], "note": note},
            "required": [*schema["required"], "note"]}


def test_generate_fields_are_independent_of_each_other(model):
    """Generate fields condition on decided/extracted values only and decode as
    batch-invariant rows: /reply is the same alone and next to another generate field."""
    for opts in (FillOptions(), FillOptions(temperature=0.9, seed=3)):
        alone = fill(model, CONTEXT, SCHEMA, opts)
        both = fill(model, CONTEXT, _with_note(), opts)
        assert both["value"]["reply"] == alone["value"]["reply"]
        assert both["meta"]["/reply"] == alone["meta"]["/reply"]
        assert "note" in both["value"]


def test_per_field_generate_budget(model):
    resp = fill(model, CONTEXT, _with_note(budget=3), FillOptions(max_new_tokens=64))
    assert resp["meta"]["/note"]["generated_tokens"] <= 3
    assert resp["meta"]["/reply"]["generated_tokens"] > 3 or "reply" in resp["value"]


def test_fill_many_matches_fill_alone(model):
    from deg.infer import fill_many

    other = "Dana again: the desk arrived broken. Order AB-9 total $80."
    jobs = [(CONTEXT, SCHEMA, None), (other, _sub(["department", "order_id", "total"]), None),
            (CONTEXT, {"type": "object", "properties": {"x": {"type": "object"}},
                       "required": ["x"], "bogus": 1}, None),
            (other, _with_note(), FillOptions(temperature=0.7, seed=5))]
    many = fill_many(model, jobs)
    assert len(many) == len(jobs)
    for job, got in zip(jobs, many):
        try:
            want = fill(model, *job)
        except Exception as exc:  # noqa: BLE001
            assert type(got) is type(exc)
            continue
        assert got == want


def test_score_level_labels_become_option_lines():
    from deg.infer import outcome_texts

    schema = {"type": "object", "properties": {"anger": {
        "type": "integer", "minimum": 0, "maximum": 2,
        "x-levels": {"0": "calm", "2": "furious"}}}, "required": ["anger"]}
    fld = compile_schema(schema).fields[0]
    assert fld.level_labels == ("calm", None, "furious")
    assert outcome_texts(fld) == ("0: calm", "1", "2: furious")


# ------------------------------------------------------------------- long context


class _LargestTensor:
    """Records the largest tensor (numel) any torch op produces while active."""

    def __enter__(self):
        from torch.utils._python_dispatch import TorchDispatchMode

        rec = self

        class Mode(TorchDispatchMode):
            def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                out = func(*args, **(kwargs or {}))
                for t in (out if isinstance(out, (tuple, list)) else (out,)):
                    if isinstance(t, torch.Tensor):
                        rec.largest = max(rec.largest, t.numel())
                return out

        self.largest = 0
        self.mode = Mode()
        self.mode.__enter__()
        return self

    def __exit__(self, *exc):
        self.mode.__exit__(*exc)


def test_long_prefill_streams_in_chunks_without_quadratic_tensors(model, monkeypatch):
    """No O(L^2) tensor anywhere in a long prefill (no dense L x L causal mask, no
    [B, L, P+L] mask), and the chunked prefill equals the one-shot read."""
    import deg.model as dm

    rend = model.renderer
    ids = rend.context(" ".join([CONTEXT] * 50)).ids
    L = len(ids)
    assert L > 2000
    x, n = torch.tensor([ids]), torch.tensor([L])
    with torch.no_grad():
        h_one, st_one = model.backbone(x, n)  # one shot (L <= DENSE_PREFILL_MAX)
        monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 256)
        monkeypatch.setattr(dm, "PREFILL_CHUNK", 256)
        with _LargestTensor() as big:
            h, st = model.backbone(x, n)
    assert big.largest < L * L // 4, big.largest  # was L*L (causal) and L*(P+L) (mask)
    torch.testing.assert_close(h, h_one, atol=2e-4, rtol=2e-4)
    for (r, t), (r1, t1) in zip(st.lin, st_one.lin):
        torch.testing.assert_close(r, r1, atol=2e-4, rtol=2e-4)
        torch.testing.assert_close(t, t1, atol=2e-4, rtol=2e-4)
    for (k, v), (k1, v1) in zip(st.kv, st_one.kv):
        torch.testing.assert_close(k, k1, atol=2e-4, rtol=2e-4)
        torch.testing.assert_close(v, v1, atol=2e-4, rtol=2e-4)
    assert torch.equal(st.kv_mask, st_one.kv_mask) and torch.equal(st.lengths, st_one.lengths)


def test_long_prefill_padded_batch_and_prefix_state(model, monkeypatch):
    """Chunked prefill over right-padded rows and on top of an existing prefix state
    equals the one-shot read for every row's real tokens."""
    import deg.model as dm

    rend = model.renderer
    a = rend.enc(" ".join([CONTEXT] * 12))
    b = rend.enc(" ".join([CONTEXT] * 5))
    pre = rend.enc("Note: ")
    x = torch.full((2, len(a)), rend.pad_id)
    x[0], x[1, :len(b)] = torch.tensor(a), torch.tensor(b)
    n = torch.tensor([len(a), len(b)])
    with torch.no_grad():
        _, st0 = model.backbone(torch.tensor([pre, pre]), torch.tensor([len(pre)] * 2))
        h_one, s_one = model.backbone(x, n, st0)
        monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 64)
        monkeypatch.setattr(dm, "PREFILL_CHUNK", 64)
        h, s = model.backbone(x, n, st0)
    for row, m in enumerate((len(a), len(b))):
        torch.testing.assert_close(h[row, :m], h_one[row, :m], atol=2e-4, rtol=2e-4)
    for (r, t), (r1, t1) in zip(s.lin, s_one.lin):
        torch.testing.assert_close(r, r1, atol=2e-4, rtol=2e-4)
        torch.testing.assert_close(t, t1, atol=2e-4, rtol=2e-4)
    assert torch.equal(s.kv_mask, s_one.kv_mask)
    for (k, v), (k1, v1) in zip(s.kv, s_one.kv):
        for row in range(2):  # real positions only (pad K/V are never attended)
            real = s.kv_mask[row]
            torch.testing.assert_close(k[row][:, real], k1[row][:, real], atol=2e-4, rtol=2e-4)
            torch.testing.assert_close(v[row][:, real], v1[row][:, real], atol=2e-4, rtol=2e-4)


def test_fork_batch_split_on_oom_gives_the_same_outputs(model, monkeypatch):
    """A fork batch that runs out of memory is split in half and retried; the outputs are
    the same bits as the unsplit batch."""
    from deg.model import Backbone

    flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")]
    want = _raw_outputs(model, _sub(flat))
    orig = Backbone.fork_forward
    calls = []

    def tight(self, ids, lengths, fs, keep=False, decode=False, sel=None):
        calls.append(ids.shape[0])
        if ids.shape[0] > 2:
            raise torch.OutOfMemoryError("synthetic")
        return orig(self, ids, lengths, fs, keep, decode, sel)

    monkeypatch.setattr(Backbone, "fork_forward", tight)
    got = _raw_outputs(model, _sub(flat))
    assert max(calls) > 2 and any(c <= 2 for c in calls)
    assert got.keys() == want.keys()
    for k in want:
        assert torch.equal(got[k], want[k]), k


def test_blocked_span_logits_match_one_shot(model, monkeypatch):
    """Span logits computed in answer-row micro-batches and context-row blocks (the path
    for long contexts / big schemas) equal the one-shot path, and big sets wait in host
    memory, not on the device."""
    import deg.model as dm

    flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")]
    want = _raw_outputs(model, _sub(flat))
    monkeypatch.setattr(dm, "HEAD_ROWS", 7)  # several context blocks
    monkeypatch.setattr(dm, "HEAD_LOGIT_BYTES", 8 * 3 * 64)  # micro-batches of few rows
    got = _raw_outputs(model, _sub(flat))
    assert got.keys() == want.keys()
    for k in want:  # CPU reference GEMMs are per call: allow last-bit noise only
        torch.testing.assert_close(got[k], want[k], atol=1e-5, rtol=1e-5)


def test_fork_hidden_keeps_only_head_positions(model):
    """Between the fork pass and the heads only the answer/option rows of each fork are
    held (not every fork token), and outputs are unchanged."""
    from deg.model import _Picked

    h = torch.arange(12.0).reshape(6, 2)
    p = _Picked(h, (5, 1, 3))
    assert p.rows.shape == (3, 2)
    for i in (1, 3, 5):
        assert torch.equal(p[i], h[i])


def test_consume_extends_state_without_holding_two_copies(model):
    """forward(consume=True) gives the same outputs as a normal extension and releases the
    given state's K/V (so a long context is never held twice); same for the chunked
    path."""
    import deg.model as dm

    rend = model.renderer
    ctx = torch.tensor([rend.enc(" ".join([CONTEXT] * 3))])
    more = torch.tensor([rend.enc("KNOWN VALUES:\n/a = 1\n")])
    with torch.no_grad():
        _, st = model.backbone(ctx, torch.tensor([ctx.shape[1]]))
        h_ref, s_ref = model.backbone(more, torch.tensor([more.shape[1]]), st)
        assert all(kv is not None for kv in st.kv)
        h, s = model.backbone(more, torch.tensor([more.shape[1]]), st, consume=True)
        assert all(kv is None for kv in st.kv)
        assert torch.equal(h, h_ref)
        for (k, v), (k1, v1) in zip(s.kv, s_ref.kv):
            assert torch.equal(k, k1) and torch.equal(v, v1)
        _, st = model.backbone(ctx, torch.tensor([ctx.shape[1]]))
        old = dm.DENSE_PREFILL_MAX, dm.PREFILL_CHUNK
        dm.DENSE_PREFILL_MAX = dm.PREFILL_CHUNK = 4
        try:
            h2, s2 = model.backbone(more, torch.tensor([more.shape[1]]), st, consume=True)
        finally:
            dm.DENSE_PREFILL_MAX, dm.PREFILL_CHUNK = old
        assert all(kv is None for kv in st.kv)
        torch.testing.assert_close(h2, h_ref, atol=2e-4, rtol=2e-4)


def test_span_candidates_band_is_linear_and_unchanged():
    """Span candidates over a long context never build the T x T grid; on a short one
    they equal the dense-grid reference."""
    from deg.infer import MAX_SPAN_TOKENS, SPAN_CANDIDATES, span_candidates

    def dense(start, end, temperature):
        ls = torch.log_softmax(start.float() / temperature, -1)
        le = torch.log_softmax(end.float() / temperature, -1)
        T = ls.shape[0] - 1
        grid = ls[:T, None] + le[None, :T]
        idx = torch.arange(T)
        ok = (idx[None] >= idx[:, None]) & (idx[None] - idx[:, None] < MAX_SPAN_TOKENS)
        grid = grid.masked_fill(~ok, float("-inf"))
        vals, flat = grid.flatten().topk(min(SPAN_CANDIDATES, int(ok.sum())))
        return [(float(v), int(f) // T, int(f) % T) for v, f in zip(vals, flat)]

    g = torch.Generator().manual_seed(0)
    temp = torch.tensor(1.3)
    for T in (1, 5, 127, 128, 129, 400):
        s, e = torch.randn(T + 1, generator=g), torch.randn(T + 1, generator=g)
        got, null = span_candidates(s, e, temp)
        assert got == dense(s, e, temp), T
    T = 20000
    s, e = torch.randn(T + 1, generator=g), torch.randn(T + 1, generator=g)
    with _LargestTensor() as big:
        got, _ = span_candidates(s, e, temp)
    assert big.largest <= T * MAX_SPAN_TOKENS + MAX_SPAN_TOKENS
    assert all(0 <= j - i < MAX_SPAN_TOKENS for _, i, j in got)


def test_token_span_bisect_matches_linear_scan():
    offsets = ((0, 3), (3, 4), (4, 4), (4, 9), (9, 12), (12, 12), (13, 20))
    for start in range(0, 20):
        for end in range(start + 1, 21):
            try:
                i = next(k for k, (_, b) in enumerate(offsets) if b > start)
                j = max(k for k, (a, _) in enumerate(offsets) if a < end)
            except (StopIteration, ValueError):
                continue
            assert token_span(offsets, start, end) == (i, j), (start, end)


# ------------------------------------------------------------------------- training


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

    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]}}
    model.meta["array_mode"] = "index"  # run-1 layout: count head + "ITEM i of"
    try:
        ex_index = deg_example(model, rec, random.Random(0), shuffle_p=0.0)
    finally:
        model.meta["array_mode"] = "anchor"
    assert sorted(it.kind for it in ex_index.items) == [
        "choice", "count", "extract", "extract", "generate", "score"]
    ex = deg_example(model, rec, random.Random(0), shuffle_p=0.0)
    kinds = sorted(it.kind for it in ex.items)
    assert kinds == ["anchor", "choice", "extract", "extract", "generate", "score"]
    oid = next(it for it in ex.items if it.ptr == "/order_id")
    lo = ex.ctx.offsets[oid.starts[0]][0]
    hi = ex.ctx.offsets[oid.ends[0]][1]
    assert CONTEXT[slice(*trim_span(CONTEXT, lo, hi))] == "AB-1234"
    jev = jev_example(model, {"uid": "j", "state": {"a": 1}, "questions": [
        {"type": "noul", "instructions": "Is a set?", "criteria": {"true": "yes"},
         "label": {"probs": [0.9, 0.1]}},
        {"type": "score", "instructions": "How big?", "criteria": ["small", "big"],
         "label": {"hard": 1}}]}, random.Random(0))
    assert jev is not None and [it.kind for it in jev.items] == ["noul", "score"]
    model.backbone.train()
    try:
        losses, stats = forward_losses(model, [ex, ex_index, jev], 4096, 16)
    finally:
        model.backbone.eval()
    assert {"decide", "extract", "count", "anchor", "lm"} <= set(losses)
    total = sum(losses[k] for k in ("decide", "extract", "count", "anchor", "lm"))
    assert torch.isfinite(total)
    total.backward()
    assert model.heads.span.q_start.weight.grad is not None
    assert model.backbone.text.layers[0].linear_attn.in_proj_qkv.weight.grad is not None
    for p in [*model.heads.parameters(), *model.backbone.parameters()]:
        p.grad = None


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


@pytest.mark.skipif(SNAPSHOT is None or not torch.cuda.is_available(),
                    reason="needs CUDA and the Qwen3.5-0.8B snapshot")
def test_gpu_bitwise_independence_real_model():
    m = load_model(None, device="cuda")
    flat = [n for n in SCHEMA["properties"] if n not in ("items", "reply")]
    merged = _raw_outputs(m, _sub(flat))
    reordered = _raw_outputs(m, _sub(list(reversed(flat))))
    for name in flat:
        alone = _raw_outputs(m, _sub([name]))
        assert torch.equal(alone[f"/{name}"], merged[f"/{name}"]), name
        assert torch.equal(alone[f"/{name}"], reordered[f"/{name}"]), name
    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:
        assert torch.equal(_raw_outputs(m, _sub([name]))[f"/{name}"], big_out[f"/{name}"])
    resp = fill(m, CONTEXT, SCHEMA)
    validate_value(compile_schema(SCHEMA), resp["value"])
    # generate rows: same text alone and next to another generate field (graph decode)
    both = fill(m, CONTEXT, _with_note())
    assert both["value"]["reply"] == resp["value"]["reply"]
    assert both["meta"]["/reply"] == resp["meta"]["/reply"]
    # cross-request batching is bitwise identical to each request alone
    from deg.infer import fill_many

    other = "Dana again: the desk arrived broken. Order AB-9 total $80."
    jobs = [(CONTEXT, SCHEMA, None), (other, _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(SNAPSHOT is None or not torch.cuda.is_available(),
                    reason="needs CUDA and the Qwen3.5-0.8B snapshot")
def test_gpu_long_prefill_chunked_matches_and_is_linear_memory(monkeypatch):
    """Real model on CUDA: the chunked prefill (flash-style fork kernel over the K/V
    buffer) matches the one-shot read within bf16 noise for padded rows, and a 32k-token
    prefill's peak memory is linear (K/V + hidden + one chunk's work, no L x L mask)."""
    import deg.model as dm

    m = load_model(None, device="cuda")
    rend = m.renderer
    a = rend.context(" ".join([CONTEXT] * 70)).ids
    b = rend.context(" ".join([CONTEXT] * 30)).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", 1024)
        monkeypatch.setattr(dm, "PREFILL_CHUNK", 1024)
        h, s = m.backbone(x, n)
    for row, L in enumerate((len(a), len(b))):
        # bf16 noise floor of the one-shot path itself (same row alone vs in this batch):
        # per-position cosine min ~0.998, mean ~0.99987
        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)
    monkeypatch.undo()  # module defaults
    ctx = rend.context(" ".join([CONTEXT] * 750))
    L = len(ctx.ids)
    assert L > 32000
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    base = torch.cuda.memory_allocated()
    with torch.no_grad():
        hid, st = m.trunk([ctx])
        torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated() - base
    kept = sum(k.numel() * k.element_size() * 2 for k, _ in st.kv) + hid.numel() * 2
    work = peak - kept  # one chunk's transient work, independent of L
    assert work < 640 * 2**20, (peak / 2**20, kept / 2**20)
    assert L * L > 640 * 2**20  # the old dense causal mask alone (L^2 bool bytes)
    # span keys are projected in HEAD_ROWS blocks: bitwise the same as one GEMM
    from deg.kernels import bi_linear

    ks = m.heads.span.k_start
    rows = hid[0, :20000].float()
    whole = bi_linear(rows, ks.weight, ks.bias)
    blocks = torch.cat([bi_linear(rows[i:i + 8192], ks.weight, ks.bias)
                        for i in range(0, rows.shape[0], 8192)])
    assert torch.equal(whole, blocks)


@pytest.mark.skipif(SNAPSHOT is None or not torch.cuda.is_available(),
                    reason="needs CUDA and the Qwen3.5-0.8B snapshot")
def test_gpu_peak_memory_does_not_scale_with_fields_times_context(monkeypatch):
    """A long context (~24k tokens) with a big schema: the fill's peak device memory above
    the context's own K/V + hidden states is bounded by one fork batch + one head
    micro-batch, independent of how many fields there are (fields x context logits wait
    in host memory; fork hidden states keep only head rows), and blocked span logits are
    bitwise the one-shot ones on CUDA."""
    import deg.model as dm

    m = load_model(None, device="cuda")
    monkeypatch.setattr(dm, "FORK_TOKENS", 2048)  # fixed fork batch: comparable peaks
    ctx = " ".join([CONTEXT] * 550)
    L = len(m.renderer.context(ctx).ids)
    assert L > 20000

    def schema(n):
        return {"type": "object", "properties": {
            f"f{i}": {"type": ["string", "null"], "description": f"value {i} of the ticket"}
            for i in range(n)}}

    def peak(n):
        fill(m, ctx, schema(n))  # warm-up (kernels, allocator)
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        base = torch.cuda.memory_allocated()
        fill(m, ctx, schema(n))
        torch.cuda.synchronize()
        return torch.cuda.max_memory_allocated() - base

    small, big = peak(20), peak(400)
    old_logits = 400 * (L + 1) * 8  # what holding every field's start/end on device cost
    assert big - small < 96 * 2**20, (small / 2**20, big / 2**20)
    assert old_logits > 64 * 2**20
    # blocked == one-shot, bitwise, on CUDA
    ectx = m.renderer.context(ctx)
    with torch.no_grad():
        hid, st = m.trunk([ectx])
        ch = hid[0, ectx.ctx_start: ectx.ctx_start + len(ectx.offsets)]
        a = torch.randn(5, hid.shape[-1], device="cuda")
        blocked = m._span_logits_blocked(a, ch)
        monkeypatch.setattr(dm, "HEAD_ROWS", 10**9)
        monkeypatch.setattr(dm, "HEAD_LOGIT_BYTES", 2**40)
        whole = m._span_logits_blocked(a, ch)
    for (s1, e1), (s2, e2) in zip(blocked, whole):
        assert torch.equal(s1, s2) and torch.equal(e1, e2)