"""Memory-aware micro-batches and out-of-memory recovery in the trainer (CPU):

* the memory estimate is monotone in context tokens, field rows and query tokens;
* an oversized record is split into field-row chunks that partition its items and fit;
* training forks select their prefix rows inside the checkpoint with the same hidden
  states and gradients as the old select-then-run path;
* a fake OOM injector: split-and-retry yields the same gradients and losses as the
  unsplit step; a minimal unit that keeps failing is skipped, the step leaves no
  gradient behind, and the optimizer state matches a run that never saw that step.
"""

from __future__ import annotations

import copy
import dataclasses
import random
from types import SimpleNamespace

import pytest

torch = pytest.importorskip("torch")

from deg.memplan import (  # noqa: E402
    GIB,
    MIB,
    MemCoef,
    MemPlanner,
    estimate,
    features,
)
from deg.train import Example, Item, cap_record_parts, split_part  # noqa: E402

# ------------------------------------------------------------------ estimator (no model)


def _ex(ctx: int, rows: int, qlen: int, uid: str = "r", gen: int = 0) -> Example:
    ctx_obj = SimpleNamespace(ids=tuple(range(ctx)), offsets=tuple((i, i + 1)
                                                                   for i in range(ctx - 8)))
    items = [Item("extract" if i % 2 else "choice", tuple(range(qlen))) for i in range(rows)]
    if gen:
        items.append(Item("generate", tuple(range(qlen + gen)), prompt_len=qlen))
    return Example("deg", uid, ctx_obj, items)  # type: ignore[arg-type]


def _est(*exs: Example, ft: int = 4096, fr: int = 16) -> float:
    return estimate(features(list(exs), ft, fr), MemCoef())


def test_estimate_monotone_in_context_rows_and_query_tokens():
    for ctx in (64, 512, 2048, 4096, 6000):
        assert _est(_ex(ctx, 20, 60)) < _est(_ex(ctx + 256, 20, 60))
    for rows in (1, 8, 16, 17, 64, 200):
        assert _est(_ex(4096, rows, 60)) < _est(_ex(4096, rows + 1, 60))
    for q in (8, 32, 100, 400):
        assert _est(_ex(4096, 20, q)) < _est(_ex(4096, 20, q + 16))
    assert _est(_ex(1024, 10, 50)) < _est(_ex(1024, 10, 50, gen=40))  # LM targets
    # a short record batched beside a long one pays the padded trunk width
    assert _est(_ex(4096, 20, 60), _ex(200, 20, 60)) > _est(_ex(200, 20, 60), _ex(200, 20, 60))
    # every coefficient scales the estimate up, never down
    base = features([_ex(2048, 30, 80, gen=20)], 4096, 16)
    c = MemCoef()
    for name in c.__dataclass_fields__:
        bumped = MemCoef(**{**c.__dict__, name: getattr(c, name) * 2 + 1})
        assert estimate(base, bumped) >= estimate(base, c), name


def test_coefficient_file(tmp_path):
    f = tmp_path / "coef.json"
    f.write_text('{"fork_tok": 1000, "base": 0}')
    c = MemCoef.load(f)
    assert c.fork_tok == 1000 and c.base == 0 and c.trunk_tok == MemCoef().trunk_tok
    f.write_text('{"nope": 1}')
    with pytest.raises(ValueError):
        MemCoef.load(f)
    f.write_text('{"base": -1}')
    with pytest.raises(ValueError):
        MemCoef.load(f)


def _planner(target_gib: float) -> MemPlanner:
    return MemPlanner(target=target_gib * GIB, static=8 * GIB, accum_extra=2 * GIB,
                      fork_tokens=4096, fork_rows=16)


def test_split_rows_partitions_items_into_fitting_chunks():
    pl = _planner(16.0)
    big = _ex(6000, 1500, 120, uid="big", gen=30)
    assert not pl.fits([big])
    assert all(pl.estimate([dataclasses.replace(big, items=[it])]) <= pl.budget(True)
               for it in big.items)  # every single row fits: so must every chunk
    chunks = pl.split_rows(big)
    assert len(chunks) > 1
    assert [it for c in chunks for it in c.items] == big.items  # order kept, nothing lost
    assert all(c.ctx is big.ctx and c.uid == "big" for c in chunks)
    assert all(pl.estimate([c]) <= pl.budget(True) for c in chunks)
    small = _ex(500, 10, 40)
    assert pl.split_rows(small) == [small]


def test_plan_packs_whole_batch_or_fitting_parts():
    pl = _planner(16.0)
    batch = [_ex(300, 12, 50, uid=f"s{i}") for i in range(6)]
    assert pl.fits(batch) and pl.plan(batch) == [batch]
    batch = [_ex(800, 30, 60, uid=f"m{i}") for i in range(12)] + [_ex(6000, 1500, 120, "big")]
    parts = pl.plan(batch)
    assert len(parts) > 1
    assert all(pl.estimate(p) <= pl.budget(True) for p in parts)
    got = [(ex.uid, id(it)) for p in parts for ex in p for it in ex.items]
    assert got == [(ex.uid, id(it)) for ex in batch for it in ex.items]


def test_cap_record_parts_bounds_a_pathological_record():
    pl = _planner(16.0)
    big = _ex(6000, 1500, 120, uid="big", gen=30)
    small = [_ex(300, 12, 50, uid=f"s{i}") for i in range(3)]
    n_chunks = len(pl.split_rows(big))
    assert n_chunks > 3
    batch = [*small, big]
    same, left = cap_record_parts(pl, batch, 0, random.Random(0))  # 0 = off
    assert same == batch and left == 0
    out, left = cap_record_parts(pl, batch, 2, random.Random(0))
    assert out[:3] == small  # records that fit are untouched
    capped = out[3]
    assert capped.ctx is big.ctx and capped.uid == "big"
    assert left == len(big.items) - len(capped.items) > 0
    assert capped.subsampled == left
    ids = [id(it) for it in big.items]  # a subset, in the record's order
    pos = [ids.index(id(it)) for it in capped.items]
    assert pos == sorted(pos)
    assert len(pl.split_rows(capped)) <= 2
    assert sum(len(p) for p in pl.plan(out) if any(e.uid == "big" for e in p)) <= 2
    # a cap above the record's chunk count leaves it whole
    out, left = cap_record_parts(pl, batch, n_chunks, random.Random(0))
    assert out[3] is big and left == 0


def test_observe_calibrates_scale_up_at_once_and_down_from_a_window():
    pl = _planner(18.0)
    ex = _ex(4096, 40, 80)
    raw = pl.estimate([ex])
    assert raw > GIB
    net = raw - pl.coef.base
    pl.observe([ex], pl.coef.base + 1.5 * net)  # under-estimate: raise at once
    assert pl.scale == pytest.approx(1.5)
    assert pl.estimate([ex]) > raw
    for _ in range(pl.min_samples):
        pl.observe([ex], pl.coef.base + 0.7 * net)
    assert pl.scale == pytest.approx(1.5)  # the 1.5 sample is still in the window
    pl.ratios.clear()
    for _ in range(pl.min_samples):
        pl.observe([ex], pl.coef.base + 0.7 * net)
    assert pl.scale == pytest.approx(0.7)
    pl.observe([_ex(16, 1, 4)], 100 * MIB)  # tiny batch: workspace noise, ignored
    assert pl.scale == pytest.approx(0.7)


def test_split_part_halves_records_then_rows():
    a, b, c = (_ex(100, 5, 10, uid=u) for u in "abc")
    assert split_part([a, b, c]) == [[a], [b, c]]
    halves = split_part([a])
    assert halves is not None and [len(h[0].items) for h in halves] == [2, 3]
    assert split_part([_ex(100, 1, 10)]) is None
    assert split_part([_ex(100, 3, 10)], min_rows=3) is None


# ---------------------------------------------------------- tiny model (CPU, real code)

pytest.importorskip("transformers")
from test_model import CONTEXT, SNAPSHOT, _sub, tiny_lm  # noqa: E402

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


@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)
    m.backbone.grad_ckpt = True
    m.backbone.text.embed_tokens.weight.requires_grad_(False)
    return m


def _examples(model) -> list[Example]:
    from deg.train import deg_example, jev_example

    a = CONTEXT.index("AB-1234")
    out = []
    for n, (dept, reply) in enumerate((("billing", "Refund issued."),
                                       ("technical", "We are on it."),
                                       ("sales", "Charged once now."))):
        ctx = CONTEXT if n != 1 else CONTEXT + " Extra note: the customer called twice."
        rec = {"uid": f"t{n}", "schema": _sub(["department", "order_id", "items", "reply"]),
               "context": ctx,
               "value": {"department": dept, "order_id": "AB-1234",
                         "items": [{"name": "chairs", "qty": 2}], "reply": reply},
               "spans": {"/order_id": [a, a + 7],
                         "/items/0/name": [CONTEXT.index("chairs"),
                                           CONTEXT.index("chairs") + 6]},
               "span_alternatives": {}, "span_probs": {}, "decide_probs": {}}
        out.append(deg_example(model, rec, random.Random(n), shuffle_p=0.0))
    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
    return [*out, jev]


def test_fork_rows_selected_inside_checkpoint_match_select_path(model):
    exs = _examples(model)[:2]
    model.backbone.train()
    try:
        outs = []
        for inside in (True, False):
            for p in model.backbone.parameters():
                p.grad = None
            _, state = model.trunk([ex.ctx for ex in exs])
            seqs = [it.seq for ex in exs for it in ex.items][:6]
            rows = torch.tensor([0, 0, 0, 1, 1, 1])
            if inside:
                hidden = model.run_forks(state, rows, seqs)
            else:  # the previous path: per-row copies of the prefix state
                width = max(len(s) for s in seqs)
                ids = torch.full((len(seqs), width), model.renderer.pad_id, dtype=torch.long)
                for i, s in enumerate(seqs):
                    ids[i, :len(s)] = torch.tensor(s)
                lengths = torch.tensor([len(s) for s in seqs])
                h, _ = model.backbone(ids, lengths, state.select(rows))
                hidden = [h[i, :len(s)] for i, s in enumerate(seqs)]
            loss = sum((h.float() ** 2).mean() for h in hidden)
            loss.backward()
            grads = {n: p.grad.clone() for n, p in model.backbone.named_parameters()
                     if p.grad is not None}
            outs.append(([h.detach() for h in hidden], grads))
    finally:
        model.backbone.eval()
    (h1, g1), (h2, g2) = outs
    for a, b in zip(h1, h2):
        torch.testing.assert_close(a, b, rtol=1e-5, atol=1e-6)
    assert g1.keys() == g2.keys() and g1
    for n in g1:
        torch.testing.assert_close(g1[n], g2[n], rtol=1e-4, atol=1e-6, msg=n)


class _Run:
    """Trainer state as in main(): fp32 masters for the backbone, AdamW over masters +
    heads."""

    def __init__(self, model) -> None:
        self.model = model
        self.bb = [p for p in model.backbone.parameters() if p.requires_grad]
        self.hd = list(model.heads.parameters())
        self.master = [p.detach().float().clone().requires_grad_(True) for p in self.bb]
        self.opt = torch.optim.AdamW([{"params": self.master, "lr": 1e-3},
                                      {"params": self.hd, "lr": 1e-3}])

    def step(self, parts, fault=None):
        from deg.train import Weights, run_step

        res = run_step(self.model, parts, Weights(), 4096, 4, self.bb, self.master, self.hd,
                       live=True, fault=fault)
        if res.ok:
            self.opt.step()
            with torch.no_grad():
                for p, mp in zip(self.bb, self.master):
                    p.copy_(mp)
        return res

    def grads(self) -> list[torch.Tensor | None]:
        return [None if q.grad is None else q.grad.clone() for q in (*self.master, *self.hd)]

    def snapshot(self) -> tuple[dict, dict, dict]:
        return (copy.deepcopy(self.model.backbone.state_dict()),
                copy.deepcopy(self.model.heads.state_dict()),
                copy.deepcopy(self.opt.state_dict()))


def _same(x, y) -> None:
    """Exact equality of nested state (tensors bitwise, everything else ==)."""
    if isinstance(x, torch.Tensor):
        assert torch.equal(x, y)
    elif isinstance(x, dict):
        assert x.keys() == y.keys()
        for k in x:
            _same(x[k], y[k])
    elif isinstance(x, (list, tuple)):
        assert len(x) == len(y)
        for a, b in zip(x, y):
            _same(a, b)
    else:
        assert x == y


def _oom(*_a) -> None:
    raise torch.OutOfMemoryError("CUDA out of memory. (injected)")


def test_oom_split_retry_gives_the_unsplit_gradient(model):
    from deg.train import Weights, run_step

    exs = _examples(model)
    bb = [p for p in model.backbone.parameters() if p.requires_grad]
    hd = list(model.heads.parameters())
    model.backbone.train()
    model.heads.train()
    try:
        master = [p.detach().float().clone().requires_grad_(True) for p in bb]
        ref = run_step(model, [exs], Weights(), 4096, 4, bb, master, hd, live=True)
        assert ref.ok and ref.parts == 1 and ref.oom_retry == 0
        g_ref = [None if q.grad is None else q.grad.clone() for q in (*master, *hd)]
        assert sum(g is not None for g in g_ref) > len(master)  # backbone + heads trained

        seen: list[tuple[str, int, int]] = []

        def fault(part, phase):  # "memory" holds 1 record with <= 3 rows at most
            seen.append((phase, len(part), sum(len(e.items) for e in part)))
            if len(part) > 1 and phase == "forward":
                _oom()
            if len(part[0].items) > 3 and phase == "backward":  # grads already written
                _oom()

        master2 = [p.detach().float().clone().requires_grad_(True) for p in bb]
        got = run_step(model, [exs], Weights(), 4096, 4, bb, master2, hd, live=True,
                       fault=fault)
    finally:
        model.backbone.eval()
        model.heads.eval()
    assert got.ok and got.oom_retry > 0 and not got.oom_skip and got.parts > len(exs)
    assert any(ph == "backward" and rows > 3 for ph, _, rows in seen)
    # same gradient up to fp32 summation order (padding, batch split): per tensor within
    # 1e-4 of its norm plus 1e-6 of the whole gradient's norm (near-zero tensors)
    full = float(torch.stack([g.norm() for g in g_ref if g is not None]).norm())
    for q, g in zip((*master2, *hd), g_ref):
        assert (q.grad is None) == (g is None)
        if g is not None:
            assert float((q.grad - g).norm()) <= 1e-4 * float(g.norm()) + 1e-6 * full
    assert got.losses.keys() == ref.losses.keys()
    for k in ref.losses:
        assert got.losses[k] == pytest.approx(ref.losses[k], rel=1e-4, abs=1e-6), k
    assert all(p.grad is None for p in bb)  # bf16-side grads never linger


def test_oom_skip_leaves_no_trace_in_optimizer_or_grads(model):
    exs = _examples(model)
    good1, good2 = [exs[0], exs[3]], [exs[2]]
    bad = [exs[1], exs[0]]
    model.backbone.train()
    model.heads.train()
    init = (copy.deepcopy(model.backbone.state_dict()), copy.deepcopy(model.heads.state_dict()))
    try:
        run_a = _Run(model)
        assert run_a.step([good1]).ok and run_a.step([good2]).ok
        after_a = run_a.snapshot()

        model.backbone.load_state_dict(init[0])
        model.heads.load_state_dict(init[1])
        run_b = _Run(model)
        assert run_b.step([good1]).ok
        before = run_b.snapshot()

        def fault(part, phase):  # record t1 never fits, even one row; grads written first
            if any(e.uid == "t1" for e in part) and phase == "backward":
                _oom()

        res = run_b.step([bad], fault=fault)
        assert not res.ok and res.oom_retry > 0
        assert [e["uid"] for e in res.oom_skip] == ["t1"]
        ev = res.oom_skip[0]
        assert ev["event"] == "oom_skip" and ev["items"] == 1 and ev["context_tokens"] > 0
        assert all(g is None for g in run_b.grads())
        assert all(p.grad is None for p in model.backbone.parameters())
        mid = run_b.snapshot()
        _same(before, mid)  # nothing moved: weights, heads, optimizer moments
        assert run_b.step([good2]).ok
        after_b = run_b.snapshot()
    finally:
        model.backbone.load_state_dict(init[0])
        model.heads.load_state_dict(init[1])
        model.backbone.eval()
        model.heads.eval()
    _same(after_a, after_b)


def test_non_oom_errors_propagate(model):
    from deg.train import Weights, run_step

    exs = _examples(model)

    def fault(part, phase):
        raise ValueError("not memory")

    with pytest.raises(ValueError):
        run_step(model, [exs[:1]], Weights(), 4096, 4, [], [], list(model.heads.parameters()),
                 live=False, fault=fault)


def test_main_loop_recovers_and_logs_oom_counters(model, tmp_path, monkeypatch):
    """main() end to end on CPU (CUDA calls stubbed): planner on, an injected OOM on every
    multi-record micro-batch and on every forward of record t1 (minimal units included),
    so steps split, retry, skip t1, redraw; the run finishes and logs the counters."""
    import json

    import deg.train as tr

    recs = []
    a = CONTEXT.index("AB-1234")
    for n in range(6):
        recs.append({"uid": f"t{n}", "split": "train",
                     "schema": _sub(["department", "order_id", "items", "reply"]),
                     "context": CONTEXT + f" Note {n}.",
                     "value": {"department": "billing", "order_id": "AB-1234",
                               "items": [{"name": "chairs", "qty": 2}], "reply": "Refunded."},
                     "spans": {"/order_id": [a, a + 7]}, "span_alternatives": {},
                     "span_probs": {}, "decide_probs": {}})
    path = tmp_path / "records.jsonl"
    path.write_text("".join(json.dumps(r) + "\n" for r in recs))

    for name, fn in {"set_device": lambda *_: None, "device_count": lambda: 1,
                     "memory_allocated": lambda *_: 0, "synchronize": lambda *_: None,
                     "max_memory_allocated": lambda *_: 0}.items():
        monkeypatch.setattr(torch.cuda, name, fn)
    monkeypatch.setattr(tr, "load_model", lambda *a, **k: model)
    monkeypatch.setattr(tr, "GpuMonitor", lambda: SimpleNamespace(take=dict, stop=lambda: None))
    real = tr.forward_losses

    def flaky(m, part, *a):
        if len(part) > 1 or any(e.uid == "t1" for e in part):
            _oom()
        return real(m, part, *a)

    monkeypatch.setattr(tr, "forward_losses", flaky)
    out = tmp_path / "run"
    init = (copy.deepcopy(model.backbone.state_dict()), copy.deepcopy(model.heads.state_dict()))
    try:
        tr.main(["--out", str(out), "--records", str(path), "--jev-frac", "0", "--steps", "4",
                 "--log-every", "2", "--max-records", "4", "--budget", "4000",
                 "--source-min-train", "1", "--mem-target-gb", "2", "--no-final-eval",
                 "--oom-max-redraws", "2", "--fork-rows", "4", "--seed", "3",
                 "--lm-weight", "2"])
    finally:
        model.backbone.load_state_dict(init[0])
        model.heads.load_state_dict(init[1])
        model.backbone.eval()
        model.heads.eval()
    rows = [json.loads(x) for x in (out / "train.log.jsonl").read_text().splitlines()]
    plan = next(r for r in rows if r.get("event") == "mem_plan")
    assert plan["target_gib"] == 2 and plan["activation_budget_gib"] > 0
    steps = [r for r in rows if "step" in r and "oom_retry" in r]
    assert [r["step"] for r in steps] == [2, 4]
    assert sum(r["oom_retry"] for r in steps) > 0
    skips = [r for r in rows if r.get("event") == "oom_skip"]
    assert skips and all(r["uid"] == "t1" and r["items"] == 1 for r in skips)
    assert sum(r["oom_skip"] for r in steps) == len(skips)
    assert any("total" in r for r in steps)  # real optimizer steps happened
    assert (out / "final").exists()
    saved = json.loads((out / "final" / "deg.json").read_text())
    assert saved["train_args"]["lm_weight"] == "2.0"


def test_chunked_trunk_estimate_is_linear_in_context():
    """Above DENSE_PREFILL_MAX the trunk is read in checkpointed chunks: no quadratic
    mask term, linear growth, and a 16k-token window with forks plans under 18 GiB beside
    the measured run-2b static (~8.5 GiB) at the default coefficients."""
    from deg.model import DENSE_PREFILL_MAX

    f = features([_ex(16384, 40, 80, gen=30)], 8192, 32)
    assert f.trunk_sq == 0 and f.trunk_chunk_tok == 16384 and f.trunk_chunks == 4
    a, b, c = (_est(_ex(w, 40, 80), ft=8192, fr=32) for w in (10000, 13000, 16000))
    assert a < b < c
    assert abs((c - b) - (b - a)) < 0.05 * (c - a)  # linear, not quadratic
    assert _est(_ex(DENSE_PREFILL_MAX, 40, 80)) == pytest.approx(
        estimate(features([_ex(DENSE_PREFILL_MAX, 40, 80)], 4096, 16, dense_max=1 << 30,
                          chunk=4096), MemCoef()))
    assert 8.5 * GIB + estimate(f, MemCoef()) < 18 * GIB


def _grad_of(model, parts, **kw):
    from deg.train import Weights, run_step

    bb = [p for p in model.backbone.parameters() if p.requires_grad]
    hd = list(model.heads.parameters())
    master = [p.detach().float().clone().requires_grad_(True) for p in bb]
    res = run_step(model, parts, Weights(anchor=0.5, anchor_count=0.5), 4096, 4, bb, master,
                   hd, live=True, **kw)
    assert res.ok
    return res, [None if q.grad is None else q.grad.clone() for q in (*master, *hd)]


def _assert_same_grads(got, ref):
    full = float(torch.stack([g.norm() for g in ref if g is not None]).norm())
    for g, r in zip(got, ref):
        assert (g is None) == (r is None)
        if r is not None:
            assert float((g - r).norm()) <= 1e-4 * float(r.norm()) + 1e-6 * full


def _chunked(ex):
    return [[dataclasses.replace(ex, items=[it])] for it in ex.items]


def test_row_chunks_share_one_trunk_pass_with_the_unsplit_gradient(model, monkeypatch):
    """A record split into row chunks runs its trunk once (forks backprop into detached
    trunk outputs, then one trunk backward) and gets the unsplit record's gradient and
    losses; chunks of different records never share a trunk."""
    exs = _examples(model)
    a, b, c, jev = exs
    assert len(a.items) >= 3 and len(c.items) >= 3
    model.backbone.train()
    model.heads.train()
    calls: list[int] = []
    trunk = model.trunk

    def counting(ctxs):
        calls.append(len(ctxs))
        return trunk(ctxs)

    try:
        ref_res, ref = _grad_of(model, [[a], [b, jev], [c]])
        monkeypatch.setattr(model, "trunk", counting)
        parts = [*_chunked(a), [b, jev], *_chunked(c)]
        got_res, got = _grad_of(model, parts)
        assert calls == [1, 2, 1]  # a once, [b, jev] once, c once
        assert got_res.trunk_shared == 2 and got_res.parts == len(parts)
        assert got_res.stats["trunk_tokens"] == ref_res.stats["trunk_tokens"]
        calls.clear()
        old_res, old = _grad_of(model, parts, share_trunk=False)
        assert len(calls) == len(parts) and old_res.trunk_shared == 0
    finally:
        monkeypatch.setattr(model, "trunk", trunk)
        model.backbone.eval()
        model.heads.eval()
    _assert_same_grads(got, ref)
    _assert_same_grads(old, ref)
    assert got_res.losses.keys() == ref_res.losses.keys()
    for k in ref_res.losses:
        assert got_res.losses[k] == pytest.approx(ref_res.losses[k], rel=1e-4, abs=1e-6), k


@pytest.mark.parametrize("phase", ["forward", "backward", "trunk"])
def test_shared_trunk_oom_falls_back_to_per_chunk_trunks(model, phase):
    """An out-of-memory error anywhere in a shared-trunk run (a later chunk's forward or
    backward, or the trunk's backward) drops the run's gradients and reruns its chunks one
    by one: same gradient as the unsplit record."""
    a = _examples(model)[0]
    chunks = _chunked(a)
    hit: list[str] = []

    def fault(part, ph):
        if ph == phase and not hit and (ph == "trunk" or part is chunks[-1]):
            hit.append(ph)
            _oom()

    model.backbone.train()
    model.heads.train()
    try:
        _, ref = _grad_of(model, [[a]])
        res, got = _grad_of(model, chunks, fault=fault)
    finally:
        model.backbone.eval()
        model.heads.eval()
    assert hit == [phase] and res.oom_retry == 1 and res.trunk_shared == 0
    assert res.parts == len(chunks)
    _assert_same_grads(got, ref)
    assert all(p.grad is None for p in model.backbone.parameters())


def test_shared_trunk_on_chunked_long_trunk_path(model, monkeypatch):
    """Same equivalence when the trunk takes the chunked checkpointed read (contexts over
    DENSE_PREFILL_MAX): its recurrent state / K/V leaves carry the forks' gradient."""
    import deg.model as dm

    monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 48)
    monkeypatch.setattr(dm, "PREFILL_CHUNK", 32)
    a = _examples(model)[0]
    assert len(a.ctx.ids) > 64
    spy: list[int] = []
    orig = dm.Backbone._forward_chunked_ckpt

    def counting(self, *args, **kw):
        spy.append(1)
        return orig(self, *args, **kw)

    monkeypatch.setattr(dm.Backbone, "_forward_chunked_ckpt", counting)
    model.backbone.train()
    model.heads.train()
    try:
        _, ref = _grad_of(model, [[a]])
        n_ref = len(spy)
        res, got = _grad_of(model, _chunked(a))
    finally:
        model.backbone.eval()
        model.heads.eval()
    assert n_ref >= 1 and len(spy) == 2 * n_ref  # chunked read, once per record
    assert res.trunk_shared == 1
    _assert_same_grads(got, ref)


def test_make_weights_takes_lm_weight_from_the_cli():
    from argparse import Namespace

    from deg.train import Weights, make_weights

    ns = Namespace(anchor_weight=0.5, anchor_count_weight=0.25, lm_weight=4.0)
    assert make_weights(ns) == Weights(anchor=0.5, anchor_count=0.25, lm=4.0)
    with pytest.raises(ValueError):
        make_weights(Namespace(anchor_weight=1.0, anchor_count_weight=0.0, lm_weight=-1.0))


def test_lm_weight_scales_only_the_lm_gradient(model):
    """run_step's gradient is linear in Weights.lm: g(3) - g(0) == 3 (g(1) - g(0)), and the
    logged total counts the lm term at its weight."""
    from deg.train import Weights, run_step

    exs = _examples(model)
    assert any(it.kind == "generate" for ex in exs for it in ex.items)
    model.backbone.train()
    model.heads.train()
    bb = [p for p in model.backbone.parameters() if p.requires_grad]
    hd = list(model.heads.parameters())
    grads, losses = {}, {}
    try:
        for w in (0.0, 1.0, 3.0):
            master = [p.detach().float().clone().requires_grad_(True) for p in bb]
            res = run_step(model, [exs], Weights(lm=w), 4096, 4, bb, master, hd, live=True)
            assert res.ok
            grads[w] = [q.grad.clone() for q in master]
            losses[w] = res.losses
    finally:
        model.backbone.eval()
        model.heads.eval()
    lm = losses[1.0]["lm"]
    assert lm > 0
    assert losses[3.0]["total"] == pytest.approx(losses[0.0]["total"] + 3 * lm, rel=1e-4)
    d1 = torch.cat([(a - b).flatten() for a, b in zip(grads[1.0], grads[0.0])])
    d3 = torch.cat([(a - b).flatten() for a, b in zip(grads[3.0], grads[0.0])])
    assert float(d1.norm()) > 0
    assert float((d3 - 3 * d1).norm()) <= 1e-3 * float(d3.norm())