"""B parallel byte streams over a (possibly multi-GB) memory-mapped corpus.

docs/BATCHED_TRAINING.md.  Bytes only: no tokenizer, no vocabulary.  The shard
files (sorted glob, same order as ``resolve_shard_paths``) are viewed as ONE
circular byte string, exactly the continuum ``StreamingPacketSource`` produces
with ``--loop-shards`` (shard ends and the epoch wrap are plain byte
transitions; in byte-tick mode the shard-end flush emits nothing).

Stream b sits at an absolute, ever-increasing position ``pos[b]`` (the corpus
byte is ``pos mod total``).  Each ``next()`` returns, for every stream, the
current byte, its 320-d byte-tick packet features and the target (next) byte,
then advances every stream by one byte.

Features are the ``Atomizer`` byte-tick features (``max_span_bytes == 1``)
computed vectorised: everything that depends only on the byte comes from a
256-row table built with the scalar ``Atomizer`` code itself; the stream-
dependent parts are the phase (``pos mod phase_period``), ``previous_length``
(0 on the first packet of a stream context, then 1), the context similarity and
the 16-value rolling context sketch (EMA 0.85, low nibble).  From a fresh
context this equals ``Atomizer.encode_bytes`` feature for feature (tested).

Resumable: ``state_dict()`` holds the corpus signature, the positions and the
per-stream feature context (exact resume).
"""

from __future__ import annotations

import math
from pathlib import Path
from typing import Sequence

import numpy as np
import torch

from .atomizer import Atomizer

# Column layout of the 320 features: hist 0:256, first 256:272, last 272:288,
# stats 288:304, context 304:320.
_STATS = 288
_COL_SIN, _COL_COS, _COL_PREV, _COL_SIM = _STATS + 10, _STATS + 11, _STATS + 12, _STATS + 13
_CTX = 304


class ByteCorpus:
    """Ordered shard files as one circular, memory-mapped byte string."""

    def __init__(self, paths: Sequence[str | Path]) -> None:
        self.paths = [Path(p).resolve() for p in paths]
        if not self.paths:
            raise ValueError("ByteCorpus needs at least one file")
        self.sizes = [p.stat().st_size for p in self.paths]
        if any(s == 0 for s in self.sizes):
            empty = [str(p) for p, s in zip(self.paths, self.sizes) if s == 0]
            raise ValueError(f"empty shard(s): {empty}")
        self.maps = [np.memmap(p, dtype=np.uint8, mode="r") for p in self.paths]
        self.starts = np.cumsum([0] + self.sizes[:-1]).astype(np.int64)
        self.total = int(sum(self.sizes))
        if self.total < 2:
            raise ValueError("corpus must hold at least 2 bytes")

    def signature(self) -> dict:
        return {"paths": [str(p) for p in self.paths], "sizes": list(self.sizes), "total": self.total}

    def read(self, pos: int, n: int) -> np.ndarray:
        """n bytes starting at absolute position ``pos`` (wraps around the corpus)."""
        out = np.empty(n, dtype=np.uint8)
        g = int(pos) % self.total
        done = 0
        while done < n:
            i = int(np.searchsorted(self.starts, g, side="right") - 1)
            off = g - int(self.starts[i])
            take = min(n - done, self.sizes[i] - off)
            out[done:done + take] = self.maps[i][off:off + take]
            done += take
            g = (g + take) % self.total
        return out


def byte_tick_feature_table(pack_mode: str = "linguistic") -> torch.Tensor:
    """[256, 320] static byte-tick features (context-free part), built with the
    scalar Atomizer: stream-dependent columns are overwritten per tick."""
    at = Atomizer(max_span_bytes=1, pack_mode=pack_mode)
    rows = []
    for b in range(256):
        at.reset()
        packets = at.encode_bytes(bytes([b]), reset=True, flush=True)
        assert len(packets) == 1
        rows.append(packets[0].features)
    return torch.stack(rows)


class BatchedByteStream:
    """B byte streams with vectorised byte-tick features.

    ``offsets``: starting absolute positions (default: evenly spaced over the
    corpus).  ``warmup``: bytes before each non-zero offset replayed through the
    feature context only (0.85^64 ≈ 3e-5, so 64 bytes ≈ a full replay).
    """

    def __init__(self, corpus: ByteCorpus, batch_size: int, *, offsets: Sequence[int] | None = None,
                 chunk: int = 4096, pack_mode: str = "linguistic", phase_period: float = 256.0,
                 warmup: int = 64, device: str | torch.device = "cpu") -> None:
        self.corpus = corpus
        self.B = int(batch_size)
        self.chunk = max(1, int(chunk))
        self.pack_mode = pack_mode
        self.phase_period = float(phase_period)
        self.device = torch.device(device)
        if offsets is None:
            offsets = [(b * corpus.total) // self.B for b in range(self.B)]
        if len(offsets) != self.B:
            raise ValueError("need one offset per stream")
        self.table = byte_tick_feature_table(pack_mode).to(self.device)
        # sin/cos of the phase, built with math.* exactly like Atomizer._make_features.
        per = int(self.phase_period)
        if float(per) != self.phase_period:
            raise ValueError("phase_period must be an integer")
        ph = [2.0 * math.pi * (p % self.phase_period) / self.phase_period for p in range(per)]
        self.sin_t = torch.tensor([math.sin(x) for x in ph], dtype=torch.float32, device=self.device)
        self.cos_t = torch.tensor([math.cos(x) for x in ph], dtype=torch.float32, device=self.device)
        self.pos = np.asarray([int(o) for o in offsets], dtype=np.int64)
        self.ctx = torch.zeros(self.B, 16, dtype=torch.float32, device=self.device)
        self.prev_len = torch.zeros(self.B, dtype=torch.float32, device=self.device)
        self._buf: torch.Tensor | None = None
        self._buf_i = 0
        self._ph = None  # device copy of pos mod phase_period (no per-tick host->device copy)
        if warmup > 0:
            self._warmup(int(warmup))

    # --------------------------------------------------------------- context
    def _warmup(self, n: int) -> None:
        start = self.pos.copy()
        for b in range(self.B):
            k = int(min(n, start[b]))
            if k <= 0:
                continue
            raw = torch.from_numpy(self.corpus.read(int(start[b]) - k, k).astype(np.int64))
            for v in raw.tolist():
                self._ctx_update_one(b, v)
        self._buf = None

    def _ctx_update_one(self, b: int, v: int) -> None:
        g = torch.zeros(16, dtype=torch.float32, device=self.device)
        g[v & 15] = 1.0
        self.ctx[b].mul_(0.85).add_(0.15 * g)
        self.prev_len[b] = 1.0

    def _refill(self) -> None:
        n = self.chunk + 1
        arr = np.stack([self.corpus.read(int(p), n) for p in self.pos])
        self._buf = torch.from_numpy(arr.astype(np.int64)).to(self.device)
        self._buf_i = 0

    # ------------------------------------------------------------------ next
    def next(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if self._buf is None or self._buf_i >= self.chunk:
            self._refill()
        i = self._buf_i
        cur = self._buf[:, i]
        tgt = self._buf[:, i + 1]
        feats = self.table.index_select(0, cur).clone()
        if self._ph is None:
            self._ph = torch.from_numpy(self.pos % int(self.phase_period)).to(self.device)
        ph = self._ph
        feats[:, _COL_SIN] = self.sin_t[ph]
        feats[:, _COL_COS] = self.cos_t[ph]
        feats[:, _COL_PREV] = self.prev_len
        low = cur & 15
        dot = self.ctx.gather(1, low.view(-1, 1)).squeeze(1).double()
        denom = self.ctx.norm(dim=1).double()
        sim = torch.where(denom > 1e-12, dot / denom.clamp_min(1e-300), torch.zeros_like(dot))
        feats[:, _COL_SIM] = ((sim + 1.0) / 2.0).float()
        feats[:, _CTX:] = self.ctx
        grouped = torch.zeros_like(self.ctx).scatter_(1, low.view(-1, 1), 1.0)
        self.ctx = self.ctx.mul(0.85).add(0.15 * grouped)
        self.prev_len = torch.ones_like(self.prev_len)
        self.pos += 1
        self._ph = (ph + 1) % int(self.phase_period)
        self._buf_i += 1
        return cur, feats, tgt

    __next__ = next

    def __iter__(self):
        return self

    # ------------------------------------------------------------ checkpoint
    def state_dict(self) -> dict:
        return {
            "corpus": self.corpus.signature(), "B": self.B, "pack_mode": self.pack_mode,
            "phase_period": self.phase_period, "pos": self.pos.tolist(),
            "ctx": self.ctx.detach().cpu().clone(), "prev_len": self.prev_len.detach().cpu().clone(),
        }

    def load_state_dict(self, sd: dict, *, strict_corpus: bool = True) -> None:
        if int(sd["B"]) != self.B:
            raise ValueError(f"stream state has B={sd['B']}, loader has B={self.B}")
        sig = self.corpus.signature()
        if strict_corpus and (sd["corpus"]["sizes"] != sig["sizes"]):
            raise ValueError("corpus changed since the checkpoint (sizes differ); pass strict_corpus=False to accept")
        self.pos = np.asarray(sd["pos"], dtype=np.int64)
        self.ctx = sd["ctx"].to(self.device).clone()
        self.prev_len = sd["prev_len"].to(self.device).clone()
        self._ph = None
        self._buf = None
        self._buf_i = 0