"""Runtime data loading: memmapped byte shards -> training batches.

A batch is a set of `ctx_len`-length windows sampled uniformly at random from
the flat token stream (nanoGPT-style). Because windows cross document
boundaries, we also compute, per position:

  * segment ids  -> which document the byte belongs to (cumulative BOS count).
    The model uses these to block cross-document attention.
  * position ids -> index WITHIN the current document, so RoPE resets at each
    document start instead of drifting across concatenated docs.

Both are cheap, fully vectorized, and ignored by the model when
`cfg.doc_attention_mask` is False (then it falls back to plain causal + arange).
"""
from __future__ import annotations

import os
from typing import Dict

import numpy as np
import torch

from .config import BOS_ID, ByteLMConfig


class ByteDataset:
    """Memmap-backed sampler for one split."""

    def __init__(self, cfg: ByteLMConfig, split: str):
        self.cfg = cfg
        self.split = split
        path = os.path.join(cfg.data_dir, f"{split}.bin")
        if not os.path.exists(path):
            raise FileNotFoundError(
                f"missing shard {path}; run data_prep.prepare_data() first"
            )
        # uint16 on disk; mmap so we never load the whole corpus into RAM.
        self.data = np.memmap(path, dtype=np.uint16, mode="r")
        self.n = self.data.shape[0]
        need = cfg.ctx_len + 1
        if self.n < need:
            raise ValueError(
                f"split {split!r} has {self.n} tokens < ctx_len+1 ({need}); "
                f"use a larger corpus or a smaller cfg.ctx_len"
            )

    def __len__(self) -> int:
        return self.n

    def _doc_aux(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Segment ids and within-doc position ids for a [B,T] input batch."""
        B, T = x.shape
        idx = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
        is_bos = x == BOS_ID
        seg_ids = torch.cumsum(is_bos.to(torch.int32), dim=1)
        # within-doc position: index minus index of the most recent BOS.
        bos_pos = torch.where(is_bos, idx, torch.full_like(idx, -1))
        last_bos = torch.cummax(bos_pos, dim=1).values.clamp_min(0)
        pos_ids = idx - last_bos
        return {"seg_ids": seg_ids, "pos_ids": pos_ids}

    def get_batch(self, device: str | torch.device = "cpu",
                  generator: torch.Generator | None = None) -> Dict[str, torch.Tensor]:
        cfg = self.cfg
        T = cfg.ctx_len
        hi = self.n - (T + 1)
        ix = torch.randint(0, hi + 1, (cfg.batch_size,), generator=generator)
        # Gather windows; cast uint16->int64 for embedding lookup.
        xb = torch.empty((cfg.batch_size, T), dtype=torch.long)
        yb = torch.empty((cfg.batch_size, T), dtype=torch.long)
        for i, start in enumerate(ix.tolist()):
            chunk = torch.from_numpy(self.data[start:start + T + 1].astype(np.int64))
            xb[i] = chunk[:-1]
            yb[i] = chunk[1:]
        batch = {"x": xb, "y": yb}
        if cfg.doc_attention_mask:
            batch.update(self._doc_aux(xb))
        # Move to device; non_blocking only helps with pinned CUDA transfers.
        non_blocking = isinstance(device, torch.device) and device.type == "cuda" or device == "cuda"
        if device not in ("cpu", torch.device("cpu")):
            batch = {k: v.to(device, non_blocking=non_blocking) for k, v in batch.items()}
        return batch