"""Token-shard data loading.

Shards are flat binary files of ``uint32`` token ids (uint16 is too small: our
vocab is 78080 > 65535). The loader draws random fixed-length windows across
shards — standard packed LM training; documents are separated by <EOT> so the
model learns boundaries even when a window straddles two documents.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

DTYPE = np.uint32


def write_tokens(path: str | Path, ids) -> int:
    arr = np.asarray(ids, dtype=DTYPE)
    arr.tofile(str(path))
    return arr.size


class ShardLoader:
    """Infinite iterator of random ``[batch, seqlen]`` int32 windows over token arrays."""

    def __init__(self, arrays, batch: int, seqlen: int, seed: int = 0):
        self.arrays = [a for a in arrays if len(a) > seqlen + 1]
        if not self.arrays:
            raise ValueError("no shard longer than seqlen+1")
        self.batch = batch
        self.seqlen = seqlen
        self.rng = np.random.default_rng(seed)
        w = np.array([len(a) for a in self.arrays], dtype=np.float64)
        self.weights = w / w.sum()

    def __iter__(self):
        T = self.seqlen
        while True:
            out = np.empty((self.batch, T), dtype=np.int32)
            for i in range(self.batch):
                s = int(self.rng.choice(len(self.arrays), p=self.weights))
                a = self.arrays[s]
                start = int(self.rng.integers(0, len(a) - T))
                out[i] = np.asarray(a[start:start + T], dtype=np.int32)
            yield out


def from_shard_dir(path: str | Path, batch: int, seqlen: int, seed: int = 0) -> ShardLoader:
    paths = sorted(Path(path).glob("*.bin"))
    if not paths:
        raise FileNotFoundError(f"no .bin shards in {path}")
    arrays = [np.memmap(p, dtype=DTYPE, mode="r") for p in paths]
    return ShardLoader(arrays, batch, seqlen, seed)


def synthetic_loader(vocab: int, batch: int, seqlen: int, seed: int = 0, period: int = 97) -> ShardLoader:
    """A learnable periodic stream (next = (cur+1) % period) so smoke tests can
    verify the loss actually drops, unlike i.i.d. random tokens."""
    period = min(period, vocab)
    stream = np.tile(np.arange(period, dtype=DTYPE), max(1, (seqlen * 50) // period + 1))
    return ShardLoader([stream], batch, seqlen, seed)