"""Read-once data loader (Stage 2).

Each file is cut into windows of seq_len+1 tokens that overlap by exactly one token
(x = first seq_len tokens, y = the same window shifted by one), so every token is a
prediction target once. Each file's windows are shuffled once with a fixed seed, and a
fixed schedule says which file every sample comes from, following the mix weights.

Everything is computed from (seed, step, slot), so a restart at any step continues
exactly where training stopped, with no saved loader state. A window is only reused
after its whole file has been read once.
"""
import numpy as np

# 524,288 tokens / 4096 = 128 sequences: the most any node count uses in one step.
# Each step owns 128 sample slots; node counts that use fewer just skip the rest.
SAMPLES_PER_STEP = 128


class ReadOnceDataset:
    def __init__(self, bin_paths, seq_len, weights, start_step, end_step, eot_token_id, seed=1234):
        assert end_step > start_step, "empty stage"
        self.seq_len = seq_len
        self.start_step = start_step
        self.eot = eot_token_id
        self.data = [np.memmap(p, dtype=np.uint32, mode="r") for p in bin_paths]
        self.n_windows = [(len(d) - 1) // seq_len for d in self.data]
        assert all(n > 0 for n in self.n_windows), "a file is shorter than one window"
        w = np.asarray(weights, dtype=np.float64)
        assert len(w) == len(self.data) and (w >= 0).all() and w.sum() > 0, "bad weights"
        w = w / w.sum()

        # each file's windows in one fixed shuffled order (int32 keeps RAM small)
        self.perm = [np.random.default_rng(seed + 1000 + f).permutation(n).astype(np.int32)
                     for f, n in enumerate(self.n_windows)]

        # which file each sample of the stage reads, and how many earlier samples used that file
        self.total = (end_step - start_step) * SAMPLES_PER_STEP
        self.src = np.random.default_rng(seed).choice(len(w), size=self.total, p=w).astype(np.int16)
        self.occ = np.empty(self.total, dtype=np.int64)
        for f in range(len(w)):
            idx = np.flatnonzero(self.src == f)
            self.occ[idx] = np.arange(len(idx))

    @staticmethod
    def slot(micro, i, rank, world, batch_size):
        """Position of one sequence inside its step's 128 slots; unique per (micro, i, rank)."""
        s = (micro * batch_size + i) * world + rank
        assert s < SAMPLES_PER_STEP, "more than 128 sequences per step: raise SAMPLES_PER_STEP"
        return s

    def window(self, step, slot):
        """(file index, window index, tokens) for one sample; tokens has seq_len+1 entries."""
        k = (step - self.start_step) * SAMPLES_PER_STEP + slot
        assert 0 <= k < self.total, f"step {step} is outside this stage"
        f = int(self.src[k])
        w = int(self.perm[f][self.occ[k] % self.n_windows[f]])
        start = w * self.seq_len
        tokens = np.asarray(self.data[f][start : start + self.seq_len + 1], dtype=np.int64)
        return f, w, tokens

    def get_numpy(self, step, micro, batch_size, rank, world):
        """One micro-batch as numpy arrays: inputs, targets, document ids."""
        x = np.empty((batch_size, self.seq_len), dtype=np.int64)
        y = np.empty((batch_size, self.seq_len), dtype=np.int64)
        seg = np.empty((batch_size, self.seq_len), dtype=np.int32)
        for i in range(batch_size):
            _, _, chunk = self.window(step, self.slot(micro, i, rank, world, batch_size))
            x[i], y[i] = chunk[:-1], chunk[1:]
            # document id goes up AFTER each end-of-text, so the EOT stays with its own document
            is_eot = (x[i] == self.eot).astype(np.int32)
            seg[i] = np.cumsum(is_eot) - is_eot
        return x, y, seg

    def get_batch(self, step, micro, batch_size, rank, world, device="cuda"):
        """Same as get_numpy, as torch tensors on `device`."""
        import torch
        out = []
        for a in self.get_numpy(step, micro, batch_size, rank, world):
            t = torch.from_numpy(a)
            if torch.device(device).type == "cuda":
                t = t.pin_memory().to(device, non_blocking=True)
            out.append(t)
        return tuple(out)