import numpy as np
import torch


class BinDataset:
    """Serves training batches from one or more memory-mapped .bin token files.

    Memory-mapping lets us read slices of a huge file without loading it into RAM.
    Each sample is a random window; input and target are the window shifted by one.
    Also returns segment IDs so attention can stay inside document boundaries.
    """

    def __init__(
        self,
        bin_paths: list[str],
        seq_len: int,
        weights: list[float] | None = None,
        eot_token_id: int = 151643,
    ):
        self.seq_len = seq_len
        self.eot_token_id = eot_token_id
        # open each source as a read-only memory map of uint32 token ids
        self.data = [np.memmap(p, dtype=np.uint32, mode="r") for p in bin_paths]
        self.lengths = [len(d) for d in self.data]

        # sampling weights = the data mix (e.g. 0.65 english, 0.20 code, 0.15 math)
        if weights is None:
            weights = [n / sum(self.lengths) for n in self.lengths]
        self.weights = np.array(weights) / np.sum(weights)

    def get_batch(self, batch_size: int, device: str = "cuda"):
        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):
            # pick a source according to the mix, then a random window inside it
            src = np.random.choice(len(self.data), p=self.weights)
            arr, n = self.data[src], self.lengths[src]
            start = np.random.randint(0, n - self.seq_len - 1)
            chunk = arr[start : start + self.seq_len + 1].astype(np.int64)
            x[i] = chunk[:-1]        # input:  tokens 0..seq_len-1
            y[i] = chunk[1:]         # target: same window shifted by one

            # segment IDs bump AFTER each EOT, so EOT stays with the doc it ends.
            # Example x[i]=[a,b,EOT,c,d,EOT,e] -> seg[i]=[0,0,0,1,1,1,2]
            is_eot = (x[i] == self.eot_token_id).astype(np.int32)
            seg[i] = np.cumsum(is_eot) - is_eot

        # pin_memory + non_blocking speeds up the CPU->GPU copy during training
        xt = torch.from_numpy(x).pin_memory().to(device, non_blocking=True)
        yt = torch.from_numpy(y).pin_memory().to(device, non_blocking=True)
        st = torch.from_numpy(seg).pin_memory().to(device, non_blocking=True)
        return xt, yt, st