"""Training memory planner: per-micro-batch peak memory from its shape (SPEC.md §12).

A micro-batch is one padded trunk forward (``B`` records, padded width ``W`` context
tokens) and fork groups (:func:`group_forks`: at most ``fork_rows`` rows and, unless a
single row is wider, ``fork_tokens`` padded tokens each), every fork row continuing its
record's trunk state of width ``W``; then the heads and the LM cross-entropy. With
gradient checkpointing (one saved hidden per layer) and the in-checkpoint prefix
selection of ``Backbone.forward(rows=...)``, its activation memory is modelled as

    persistent  = trunk_tok * B*W + trunk_row * B*K + trunk_mask * B*W^2 [dense]
                + trunk_chunk_tok * B*W [chunked]
                + fork_tok * sum_i l_i + fork_mask * sum_i l_i*(W+l_i)
                + head_pair * sum_items T_item + lm_saved * lm_targets
    transient   = max(trunk_layer * B*Wc + trunk_row_layer * B + attn_pair * B*Wc*W,
                      fork_layer * G_tok + fork_prefix * G_rows*(W+l_max)
                      + fork_row_layer * G_rows + attn_pair * G_tok*(W+l_max),
                      lm_peak * lm_targets)
    activation  = base + scale * (persistent + transient)

over the fork rows ``i`` of query length ``l_i`` (``n`` rows, longest ``l_max``). A trunk
of ``W <= DENSE_PREFILL_MAX`` tokens is read in one shot (``K = 1``, ``Wc = W``, its
``[B, W, W]`` mask is a checkpoint input); a longer one (training) is read in ``K =
ceil(W / PREFILL_CHUNK)`` checkpointed chunks (``Backbone._forward_chunked_ckpt``):
recurrent state saved per chunk, the chunks' K/V concatenated once more at the end
(``trunk_chunk_tok``), no saved mask, and one layer's recompute covers one chunk
(``Wc = PREFILL_CHUNK``) against all ``W`` keys. The largest fork group is bounded by
``G_rows = min(n, fork_rows)`` rows and ``G_tok = min(n*l_max, max(fork_tokens, l_max))``
padded tokens (``transient``: one layer's
recompute + backward, or the LM logits; they are never live together). Fork sums use
real lengths (the ~25% group padding is in the coefficients), and the group bound does
not depend on how the rows happen to group, so with every coefficient >= 0 the estimate
is monotone in the number of field rows and every query's tokens, and in the context
tokens within each trunk path (it drops where the chunked path takes over).

The default coefficients are analytic for Qwen3.5-0.8B in bf16 (24 layers: 18 gated
delta-net + 6 GQA attention with 2 K/V heads x 256, hidden 1024, MLP 3584, vocab
248,320; SDPA runs the memory-efficient kernel, so attention is quadratic only through
its boolean mask). ``scale`` is the calibration: the trainer measures each micro-batch's
real peak on the device and :meth:`MemPlanner.observe` keeps ``scale`` at the largest
measured/estimated ratio over a recent window (a measured linear correction: raised at
once by any under-estimate, lowered only from :attr:`MemPlanner.min_samples` samples,
clamped to [:attr:`MemPlanner.min_scale`, :attr:`MemPlanner.max_scale`]).
``--mem-coef FILE`` overrides coefficients with a JSON object of the same names (e.g. a
least-squares refit of the trainer's ``mem.jsonl``: features + measured peak per part).

The device budget for activations is ``mem_target - static`` where ``static`` is the
resident model + optimizer state + backbone gradients (bf16, plus the fp32 accumulators
when a step runs several parts), measured on the device once it exists.
"""

from __future__ import annotations

import dataclasses
import json
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Protocol, Sequence

KIB = 1024
MIB = 1024 * KIB
GIB = 1024 * MIB


class _ItemLike(Protocol):
    kind: str
    seq: tuple[int, ...]
    prompt_len: int


def group_forks(forks: list[tuple[int, Any]], fork_tokens: int, fork_rows: int
                ) -> list[list[tuple[int, Any]]]:
    """Fork groups of one micro-batch, as the trainer runs them: sorted by length; a group
    closes at the token/row budget or when padding would exceed ~25%."""
    forks = sorted(forks, key=lambda x: len(x[1].seq))
    groups: list[list[tuple[int, Any]]] = []
    cur: list[tuple[int, Any]] = []
    for f in forks:
        width = len(f[1].seq)
        if cur and ((len(cur) + 1) * width > fork_tokens or len(cur) >= fork_rows
                    or width > 1.25 * len(cur[0][1].seq) + 8):
            groups.append(cur)
            cur = []
        cur.append(f)
    if cur:
        groups.append(cur)
    return groups


@dataclass(frozen=True)
class MemCoef:
    """Bytes per unit of each feature (see the module docstring)."""

    base: float = 512 * MIB  # per micro-batch: cuBLAS/compile workspaces, fragmentation
    # persistent until the backward reaches them
    trunk_tok: float = 84 * KIB  # saved hidden per layer (48K) + K/V (12K) + its grad (12K)
    #                               + final hidden / head projections (~12K)
    trunk_row: float = 36 * MIB  # recurrent state [16,128,128] fp32 x 18 layers + its grad
    #                               (per row and per chunk of a chunked trunk)
    trunk_chunk_tok: float = 24 * KIB  # chunked trunk: concatenated K/V copy + its grad
    trunk_mask: float = 1.0  # [B, W, W] bool causal mask (checkpoint input)
    fork_tok: float = 70 * KIB  # saved hidden per layer + final hidden (x1.25 padding)
    fork_mask: float = 1.25  # [R, L, W+L] bool mask per group (checkpoint input)
    head_pair: float = 48.0  # span/anchor logits per (item, context token), fp32 + grads
    lm_saved: float = 4.0 * 248_320  # log-softmax output kept for the backward (fp32)
    # transient: one layer recomputed with grad + its backward
    trunk_layer: float = 352 * KIB  # projections, conv, delta-rule chunks, MLP, + grads
    trunk_row_layer: float = 8 * MIB  # one layer's recurrent state copies + grads
    fork_layer: float = 352 * KIB
    fork_prefix: float = 8 * KIB  # selected + concatenated prefix K/V and their grads
    fork_row_layer: float = 8 * MIB
    attn_pair: float = 6.0  # bool mask -> additive bias for the efficient SDPA kernel
    lm_peak: float = 14.0 * 248_320  # bf16 logits + fp32 copy + log-softmax + grads

    @classmethod
    def for_backbone(cls, backbone: str | None = None) -> MemCoef:
        """Analytic defaults per backbone (``deg.model.BACKBONES``): the class defaults
        are Qwen3.5-0.8B; ``qwen3.5-2b`` scales its width terms; ``lfm2.5-1.2b`` (16 layers: 10 short-conv + 6 GQA attention with
        8 K/V heads x 64, hidden 2048, MLP 8192, vocab 65,536) saves a 2048-wide hidden
        per layer (64K/token over 16 layers vs 48K), the same 12K/token of K/V, no
        recurrent state (conv tails only: 2 tokens x 2048 per conv layer), a ~2.3x wider
        layer recompute (hidden 2x, MLP 8192 vs 3584) and a 3.8x smaller vocabulary. The
        trainer's measured ``scale`` corrects the remaining error either way."""
        if backbone in (None, "qwen3.5-0.8b"):
            return cls()
        if backbone == "qwen3.5-2b":
            # same 24-layer layout, recurrent state and K/V as the 0.8B; hidden 2048 (2x)
            # doubles the saved hidden per layer (96K/token) and MLP 6144 (1.7x) with the
            # 2x hidden makes a layer's recompute ~1.9x wider
            return cls(trunk_tok=144 * KIB, fork_tok=120 * KIB, trunk_layer=680 * KIB,
                       fork_layer=680 * KIB)
        if backbone in ("gemma-4-e2b-it", "gemma-4-e2b"):
            # Gemma 4 E2B (35 attention layers, hidden 1536, MLP 6144 / 12288 in the 20
            # K/V-sharing layers, 15 K/V slots: 3 global x 1 head x 512 + 12 sliding x 256,
            # PLE 35 x 256 per token, vocab 262,144), derived per token in bf16: saved
            # hidden per layer 35 x 3K = 105K, the PLE inputs (17.5K) and the projected PLE
            # before its norm (17.5K), global K/V 6K + the sharing source's sliding K/V 1K
            # and their grads, final hidden / heads ~12K -> ~172K; no recurrent state, a
            # sliding tail per row of 12 slots x 2 x 511 x 256 (6 MiB) + grad; one layer's
            # recompute + backward: GeGLU at 12288 (gate, up, gelu, product + grads, ~200K),
            # q at 8 x 512, norms and residuals at 1536 (~10 x 3K), PLE gate (~2K) ->
            # ~640K/token with allocator slack; the logits carry the soft-cap tanh (fp32)
            # beside the log-softmax. The measured ``scale`` corrects the rest.
            V = 262_144
            return cls(trunk_tok=172 * KIB, trunk_row=12 * MIB, trunk_chunk_tok=12 * KIB,
                       fork_tok=150 * KIB, fork_prefix=6 * KIB, lm_saved=8.0 * V,
                       trunk_layer=640 * KIB, trunk_row_layer=1 * MIB, fork_layer=640 * KIB,
                       fork_row_layer=1 * MIB, lm_peak=22.0 * V,
                       attn_pair=30.0)  # the A100 smoke ran (and calibrated, scale 1.54) at 30
        if backbone == "lfm2.5-1.2b":
            return cls(trunk_tok=112 * KIB, trunk_row=1 * MIB, fork_tok=96 * KIB,
                       lm_saved=4.0 * 65_536, trunk_layer=800 * KIB, trunk_row_layer=1 * MIB,
                       fork_layer=800 * KIB, fork_row_layer=1 * MIB, lm_peak=14.0 * 65_536)
        raise ValueError(f"no memory coefficients for backbone {backbone!r}")

    @classmethod
    def load(cls, path: Path | str | None, backbone: str | None = None) -> MemCoef:
        """The backbone's defaults (:meth:`for_backbone`), overridden by the JSON at
        ``path`` when given."""
        base = cls.for_backbone(backbone)
        if path is None:
            return base
        data = json.loads(Path(path).read_text())
        names = {f.name for f in dataclasses.fields(cls)}
        bad = sorted(set(data) - names)
        if bad:
            raise ValueError(f"unknown memory coefficients {bad} in {path}")
        if any(float(v) < 0 for v in data.values()):
            raise ValueError(f"memory coefficients must be >= 0: {path}")
        return dataclasses.replace(base, **{k: float(v) for k, v in data.items()})


@dataclass(frozen=True)
class MemFeatures:
    trunk_rows: int = 0  # B
    width: int = 0  # W
    trunk_tok: int = 0  # B * W (padded)
    trunk_sq: int = 0  # B * W^2 (dense trunk mask; 0 for a chunked trunk)
    trunk_chunks: int = 0  # B * K (chunks per row; K = 1 dense)
    trunk_chunk_tok: int = 0  # B * W for a chunked trunk, else 0
    trunk_layer_tok: int = 0  # B * Wc: tokens of one layer's recompute
    trunk_attn: int = 0  # B * Wc * W: one layer's attention pairs
    fork_rows: int = 0  # n
    fork_tok: int = 0  # sum_i l_i
    fork_sq: int = 0  # sum_i l_i * (W + l_i)
    fork_lmax: int = 0  # l_max
    group_rows: int = 0  # G_rows
    group_tok: int = 0  # G_tok
    head_pairs: int = 0  # sum over extract/anchor items of their record's context tokens
    lm_targets: int = 0

    def as_dict(self) -> dict[str, Any]:
        return dataclasses.asdict(self)


def features(examples: Sequence[Any], fork_tokens: int, fork_rows: int,
             dense_max: int | None = None, chunk: int | None = None) -> MemFeatures:
    """Shape features of one micro-batch of :class:`deg.train.Example` (anything with
    ``ctx.ids``, ``ctx.offsets`` and ``items``). ``dense_max`` / ``chunk`` default to
    ``deg.model.DENSE_PREFILL_MAX`` / ``PREFILL_CHUNK`` (read at call time)."""
    if not examples:
        return MemFeatures()
    if dense_max is None or chunk is None:
        from deg import model as dm

        dense_max = dm.DENSE_PREFILL_MAX if dense_max is None else dense_max
        chunk = dm.PREFILL_CHUNK if chunk is None else chunk
    B = len(examples)
    W = max(len(ex.ctx.ids) for ex in examples)
    dense = W <= dense_max
    K = 1 if dense else -(-W // chunk)
    Wc = W if dense else chunk
    lens = [len(it.seq) for ex in examples for it in ex.items]
    n, lmax = len(lens), max(lens, default=0)
    lm = sum(len(it.seq) - it.prompt_len for ex in examples for it in ex.items
             if it.kind == "generate")
    pairs = sum(len(ex.ctx.offsets) for ex in examples for it in ex.items
                if it.kind in ("extract", "anchor"))
    return MemFeatures(B, W, B * W, B * W * W if dense else 0, B * K, 0 if dense else B * W,
                       B * Wc, B * Wc * W, n, sum(lens), sum(x * (W + x) for x in lens),
                       lmax, min(n, fork_rows), min(n * lmax, max(fork_tokens, lmax)),
                       pairs, lm)


def estimate(f: MemFeatures, c: MemCoef, scale: float = 1.0) -> float:
    """Activation bytes of a micro-batch with features ``f`` (module docstring)."""
    if not f.trunk_rows:
        return 0.0
    W = f.width
    persistent = (c.trunk_tok * f.trunk_tok + c.trunk_row * f.trunk_chunks
                  + c.trunk_mask * f.trunk_sq + c.trunk_chunk_tok * f.trunk_chunk_tok
                  + c.fork_tok * f.fork_tok
                  + c.fork_mask * f.fork_sq + c.head_pair * f.head_pairs
                  + c.lm_saved * f.lm_targets)
    trunk = (c.trunk_layer * f.trunk_layer_tok + c.trunk_row_layer * f.trunk_rows
             + c.attn_pair * f.trunk_attn)
    fork = (c.fork_layer * f.group_tok + c.fork_prefix * f.group_rows * (W + f.fork_lmax)
            + c.fork_row_layer * f.group_rows + c.attn_pair * f.group_tok * (W + f.fork_lmax))
    transient = max(trunk, fork, c.lm_peak * f.lm_targets)
    return c.base + scale * (persistent + transient)


def static_bytes(trainable_backbone: int, heads: int, weights: int,
                 accumulate: bool) -> float:
    """Resident bytes outside the micro-batch activations: bf16 weights (``weights``
    bytes as stored), fp32 master + AdamW moments of the trainable backbone and of the
    heads, the backbone's bf16 gradients while a backward runs, and (``accumulate``,
    several parts per step) the fp32 gradient accumulators."""
    return (weights + 12.0 * trainable_backbone + 2.0 * trainable_backbone
            + 12.0 * heads + (4.0 * (trainable_backbone + heads) if accumulate else 0.0))


@dataclass
class MemPlanner:
    """Caps micro-batches to ``target`` bytes of device memory.

    ``static`` = resident bytes outside activations for a one-part step,
    ``accum_extra`` = the extra fp32 gradient accumulators a multi-part step keeps."""

    target: float
    static: float
    accum_extra: float
    fork_tokens: int
    fork_rows: int
    coef: MemCoef = field(default_factory=MemCoef)
    scale: float = 1.0
    min_scale: float = 0.5
    max_scale: float = 4.0
    window: int = 200
    min_samples: int = 20
    ratios: deque = field(default_factory=deque)

    def budget(self, accumulate: bool = False) -> float:
        return self.target - self.static - (self.accum_extra if accumulate else 0.0)

    def estimate(self, examples: Sequence[Any]) -> float:
        return estimate(features(examples, self.fork_tokens, self.fork_rows), self.coef,
                        self.scale)

    def fits(self, examples: Sequence[Any], accumulate: bool = False) -> bool:
        return self.estimate(examples) <= self.budget(accumulate)

    def observe(self, examples: Sequence[Any], measured: float, min_bytes: float = GIB
                ) -> None:
        """Calibrate: ``measured`` = the micro-batch's activation peak on the device
        (peak allocated minus allocated before it). Micro-batches estimated under
        ``min_bytes`` are dominated by workspace noise and are not used."""
        raw = estimate(features(examples, self.fork_tokens, self.fork_rows), self.coef)
        net = raw - self.coef.base
        if raw < min_bytes or net <= 0:
            return
        self.ratios.append(max(0.0, measured - self.coef.base) / net)
        while len(self.ratios) > self.window:
            self.ratios.popleft()
        top = max(self.ratios)
        if top > self.scale or len(self.ratios) >= self.min_samples:
            # raise at once on an under-estimate; lower only from a full-enough window
            self.scale = min(self.max_scale, max(self.min_scale, top))

    def split_rows(self, ex: Any, accumulate: bool = True) -> list[Any]:
        """``ex`` alone if it fits, else copies of it over consecutive chunks of its items
        (same context object; the trainer runs them on one trunk pass, ``run_chunks``),
        each chunk fitting the multi-part budget when one item can. Chunks partition the
        items in order."""
        if self.fits([ex]) or len(ex.items) <= 1:
            return [ex]
        lim = self.budget(accumulate)
        out: list[Any] = []
        cur: list[Any] = []
        for it in ex.items:
            if cur and self.estimate([dataclasses.replace(ex, items=cur + [it])]) > lim:
                out.append(dataclasses.replace(ex, items=cur))
                cur = []
            cur.append(it)
        if cur:
            out.append(dataclasses.replace(ex, items=cur))
        return out

    def plan(self, batch: Sequence[Any]) -> list[list[Any]]:
        """A batch -> the micro-batches (parts) of one optimizer step: one part when the
        whole batch fits; else records packed greedily in order into parts that fit the
        multi-part budget, a record that does not fit alone split by rows
        (:meth:`split_rows`)."""
        batch = list(batch)
        if not batch or self.fits(batch):
            return [batch]
        lim = self.budget(True)
        parts: list[list[Any]] = []
        cur: list[Any] = []
        for ex in batch:
            if cur and self.estimate(cur + [ex]) <= lim:
                cur.append(ex)
                continue
            if cur:
                parts.append(cur)
                cur = []
            if self.estimate([ex]) <= lim:
                cur = [ex]
            else:
                parts.extend([chunk] for chunk in self.split_rows(ex))
        if cur:
            parts.append(cur)
        return parts