src/pns/train/loader.py

"""Training/eval loaders.

LifetimeBatcher: batches of whole lifetimes for the streaming PNSR trainer
(TBPTT chunks are cut by the trainer; the batcher provides full arrays).
WindowSampler: (window, position) samples for the transcript baselines.

Both materialize the live cache map from the packed per-event deltas - the
same artifact certified at generation time (no independent replay logic).
"""
from __future__ import annotations

import threading
from queue import Queue

import numpy as np

from ..data.view import Shard, shard_paths
from ..world.reducer import N_PTR_SLOTS

T_EV = 48
VAL_T, KEY_T = 12, 14


def _lifetime_arrays(sh: Shard, i: int, with_meta: bool = False) -> dict:  # noqa: C901
    lo, hi = int(sh.lt_off[i]), int(sh.lt_off[i + 1])
    rlo, rhi = int(sh.rec_lt_off[i]), int(sh.rec_lt_off[i + 1])
    L = hi - lo
    tok = np.zeros((L, T_EV), np.uint16)
    for e in range(lo, hi):
        ids = sh.tokens[sh.ev_tok_off[e]:sh.ev_tok_off[e + 1]][:T_EV]
        tok[e - lo, :len(ids)] = ids
    live = np.full((L, N_PTR_SLOTS), -1, np.int16)
    cur = np.full(N_PTR_SLOTS, -1, np.int16)
    for e in range(lo, hi):
        for rid in sh.evict_rid[sh.evict_off[e]:sh.evict_off[e + 1]]:
            cur[sh.rec_slot[rlo + int(rid)]] = -1
        for rid in sh.add_rid[sh.add_off[e]:sh.add_off[e + 1]]:
            cur[sh.rec_slot[rlo + int(rid)]] = rid          # rid = local row
        live[e - lo] = cur
    out = dict(
        tok=tok, live=live,
        etype=sh.etype[lo:hi], dt=sh.dt[lo:hi],
        mode_gold=sh.mode_gold[lo:hi], family=sh.family[lo:hi],
        enum_gold=sh.enum_gold[lo:hi], enum_legal=sh.enum_legal[lo:hi],
        ptr_gold_slot=sh.ptr_gold_slot[lo:hi], op_gold=sh.op_gold[lo:hi],
        op_arg_slots=sh.op_arg_slots[lo:hi],
        rec_val_toks=sh.rec_val_toks[rlo:rhi], rec_key_toks=sh.rec_key_toks[rlo:rhi],
        rec_store=sh.rec_store[rlo:rhi], rec_kind=sh.rec_kind[rlo:rhi],
        rec_key=sh.rec_key[rlo:rhi], rec_ent=sh.rec_ent[rlo:rhi],
        rec_birth=sh.rec_birth_ev[rlo:rhi] - lo,             # lifetime-local
        rec_slot=sh.rec_slot[rlo:rhi],
    )
    # binding addresses exist only in corpora packed after E3 (the original
    # corpus predates them); consumers that need them require an E3 split.
    if hasattr(sh, "bind_write"):
        out.update(bind_write=sh.bind_write[lo:hi], bind_read=sh.bind_read[lo:hi],
                   bind_slot_ent=sh.bind_slot_ent[i],
                   bind_slot_attr=sh.bind_slot_attr[i])
    if with_meta:
        out.update(delay=sh.delay[lo:hi], kappa=sh.kappa[lo:hi],
                   corrected=sh.corrected[lo:hi], reverted=sh.reverted[lo:hi],
                   gold_val_hash=sh.gold_val_hash[lo:hi],
                   rec_val_hash=sh.rec_val_hash[rlo:rhi], seed=sh.lt_seed[i])
    return out


def _collate(lts: list[dict], with_meta: bool = False) -> dict:
    B = len(lts)
    L = max(x["etype"].shape[0] for x in lts)
    R = max(x["rec_store"].shape[0] for x in lts)
    assert all(x["etype"].shape[0] == L for x in lts), "homogeneous L required"
    b = {}
    b["tok"] = np.stack([x["tok"] for x in lts])
    b["live"] = np.stack([x["live"] for x in lts])
    for k in ("etype", "dt", "mode_gold", "family", "enum_gold", "ptr_gold_slot",
              "op_gold"):
        b[k] = np.stack([x[k] for x in lts])
    for k in ("bind_write", "bind_read", "bind_slot_ent", "bind_slot_attr"):
        if k in lts[0]:
            b[k] = np.stack([x[k] for x in lts])
    b["enum_legal"] = np.stack([x["enum_legal"] for x in lts])
    b["op_arg_slots"] = np.stack([x["op_arg_slots"] for x in lts])

    def padded(key, dtype, extra=()):
        out = np.zeros((B, R) + extra, dtype)
        for j, x in enumerate(lts):
            r = x[key].shape[0]
            out[j, :r] = x[key]
        return out
    b["rec_val_toks"] = padded("rec_val_toks", np.uint16, (VAL_T,))
    b["rec_key_toks"] = padded("rec_key_toks", np.uint16, (KEY_T,))
    for k in ("rec_store", "rec_kind", "rec_key", "rec_ent"):
        b[k] = padded(k, np.int16)
    b["rec_birth"] = padded("rec_birth", np.int32)
    b["rec_slot"] = padded("rec_slot", np.int16)
    if with_meta:
        for k in ("delay", "kappa", "corrected", "reverted"):
            b[k] = np.stack([x[k] for x in lts])
        b["gold_val_hash"] = np.stack([x["gold_val_hash"] for x in lts])
        b["rec_val_hash"] = padded("rec_val_hash", np.uint64)
        b["seed"] = np.array([x["seed"] for x in lts])
    return b


class LifetimeBatcher:
    def __init__(self, split: str, root, batch: int, seed: int = 0,
                 with_meta: bool = False, epochs: int = 1000, prefetch: int = 2,
                 max_lifetimes: int | None = None):
        self.paths = shard_paths(split, root)
        assert self.paths, f"no shards for {split} under {root}"
        self.batch, self.seed, self.with_meta = batch, seed, with_meta
        self.epochs = epochs
        self.max_lifetimes = max_lifetimes
        self.q: Queue = Queue(maxsize=prefetch)
        self._t = threading.Thread(target=self._work, daemon=True)
        self._t.start()

    def _work(self):
        rng = np.random.default_rng(self.seed)
        counts = [Shard(p).n_lifetimes for p in self.paths]
        index = [(pi, li) for pi, c in enumerate(counts) for li in range(c)]
        if self.max_lifetimes:
            index = index[: self.max_lifetimes]
        for _ in range(self.epochs):
            order = rng.permutation(len(index))
            cache: dict[int, Shard] = {}
            for k in range(0, len(order) - self.batch + 1, self.batch):
                lts = []
                for j in order[k:k + self.batch]:
                    pi, li = index[j]
                    if pi not in cache:
                        cache.clear()
                        cache[pi] = Shard(self.paths[pi])
                    lts.append(_lifetime_arrays(cache[pi], li, self.with_meta))
                self.q.put(_collate(lts, self.with_meta))
        self.q.put(None)

    def __iter__(self):
        while True:
            item = self.q.get()
            if item is None:
                return
            yield item


def iter_eval_batches(split: str, root, batch: int, limit: int | None = None):
    """Deterministic order, with meta, no shuffling; loads shard by shard."""
    n = 0
    for p in shard_paths(split, root):
        sh = Shard(p)
        for k in range(0, sh.n_lifetimes, batch):
            idx = range(k, min(k + batch, sh.n_lifetimes))
            yield _collate([_lifetime_arrays(sh, i, True) for i in idx], True)
            n += len(idx)
            if limit and n >= limit:
                return


class WindowSampler:
    """Supervised (window, position) samples for TX models."""

    def __init__(self, split, root, batch, window, seed=0, noout_frac=0.15,
                 epochs=1000, prefetch=2):
        self.paths = shard_paths(split, root)
        self.batch, self.window, self.noout_frac = batch, window, noout_frac
        self.seed, self.epochs = seed, epochs
        self.q: Queue = Queue(maxsize=prefetch)
        threading.Thread(target=self._work, daemon=True).start()

    def _work(self):
        rng = np.random.default_rng(self.seed + 7)
        for _ in range(self.epochs):
            for p in rng.permutation(self.paths):
                sh = Shard(p)
                rows = []
                for i in range(sh.n_lifetimes):
                    lo, hi = int(sh.lt_off[i]), int(sh.lt_off[i + 1])
                    sup = np.nonzero(sh.mode_gold[lo:hi] > 0)[0]
                    non = np.nonzero(sh.mode_gold[lo:hi] == 0)[0]
                    keep = non[rng.random(len(non)) < self.noout_frac]
                    for e in np.concatenate([sup, keep]):
                        rows.append((i, int(e)))
                rng.shuffle(rows)
                for k in range(0, len(rows) - self.batch + 1, self.batch):
                    chunk = sorted(rows[k:k + self.batch])
                    self.q.put(self._build(sh, chunk))
        self.q.put(None)

    def _build(self, sh, chunk):
        W = self.window
        toks = np.zeros((len(chunk), W), np.uint16)
        lts: dict[int, dict] = {}
        picks = []
        for j, (i, e) in enumerate(chunk):
            if i not in lts:
                if len(lts) > 64:
                    lts.clear()
                lts[i] = _lifetime_arrays(sh, i, with_meta=True)
            lt = lts[i]
            lo = int(sh.lt_off[i])
            t_end = int(sh.ev_tok_off[lo + e + 1])
            t_start = max(int(sh.ev_tok_off[lo]), t_end - W)
            seg = sh.tokens[t_start:t_end]
            toks[j, W - len(seg):] = seg
            picks.append((lt, e))
        b = dict(tok=toks)
        for key, dt_ in (("etype", np.uint8), ("mode_gold", np.uint8),
                         ("family", np.int8), ("enum_gold", np.int16),
                         ("ptr_gold_slot", np.int16), ("op_gold", np.int8),
                         ("delay", np.int32), ("corrected", np.int8),
                         ("reverted", np.int8)):
            b[key] = np.array([lt[key][e] for lt, e in picks], dt_)
        b["enum_legal"] = np.stack([lt["enum_legal"][e] for lt, e in picks])
        b["op_arg_slots"] = np.stack([lt["op_arg_slots"][e] for lt, e in picks])
        b["gold_val_hash"] = np.array([lt["gold_val_hash"][e] for lt, e in picks],
                                      np.uint64)
        b["lt_seed"] = np.array([int(lt["seed"]) for lt, _ in picks], np.int64)
        # per-sample live bank (variable R -> pad)
        R = max(lt["rec_store"].shape[0] for lt, _ in picks)
        for key, dt_, extra in (("rec_val_toks", np.uint16, (VAL_T,)),
                                ("rec_key_toks", np.uint16, (KEY_T,)),
                                ("rec_store", np.int16, ()), ("rec_kind", np.int16, ()),
                                ("rec_key", np.int16, ()), ("rec_ent", np.int16, ()),
                                ("rec_birth", np.int32, ()), ("rec_slot", np.int16, ()),
                                ("rec_val_hash", np.uint64, ())):
            arr = np.zeros((len(picks), R) + extra, dt_)
            for j, (lt, _) in enumerate(picks):
                r = lt[key].shape[0]
                arr[j, :r] = lt[key]
            b[key] = arr
        b["live"] = np.stack([lt["live"][e] for lt, e in picks])
        b["ev_idx"] = np.array([e for _, e in picks], np.int32)
        return b

    def __iter__(self):
        while True:
            item = self.q.get()
            if item is None:
                return
            yield item