from __future__ import annotations

import numpy as np
import torch
from torch.utils.data import IterableDataset, get_worker_info


class MemmapTokenDataset(IterableDataset):
    """Minimal contiguous adapter for an existing pretokenized stream.

    Each item contains ``seq_len`` inputs and their immediately following
    next-token labels. Consecutive items advance by ``seq_len`` tokens, so the
    transition at each sequence boundary is preserved instead of skipped.

    Formal experiments should still use the project's canonical logical-stream
    loader so exact shard order, cursor and hashes match prior runs.
    """

    def __init__(self, path: str, seq_len: int, dtype: str = "uint16", seed: int = 1337, start_sequence: int = 0):
        self.path = path
        self.seq_len = seq_len
        self.dtype = np.dtype(dtype)
        self.seed = seed
        self.start_sequence = start_sequence

    def __iter__(self):
        arr = np.memmap(self.path, mode="r", dtype=self.dtype)
        n = (len(arr) - 1) // self.seq_len
        worker = get_worker_info()
        worker_id = 0 if worker is None else worker.id
        workers = 1 if worker is None else worker.num_workers
        first = self.start_sequence + worker_id
        for i in range(first, n, workers):
            start = i * self.seq_len
            x = np.asarray(arr[start : start + self.seq_len], dtype=np.int64)
            y = np.asarray(arr[start + 1 : start + self.seq_len + 1], dtype=np.int64)
            yield torch.from_numpy(x.copy()), torch.from_numpy(y.copy())