"""
Binary token dataset.

Input:
    train.bin or val.bin produced by encode_corpus.py

Output:
    x: input token IDs [block_size]
    y: target token IDs [block_size]
"""

import json
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import Dataset


class TokenBinDataset(Dataset):
    """
    Memory-mapped token dataset.

    This avoids loading the full corpus into RAM.
    """

    def __init__(
        self,
        bin_path: str,
        meta_path: str,
        block_size: int,
    ):
        self.bin_path = Path(bin_path)
        self.meta_path = Path(meta_path)
        self.block_size = block_size

        with open(self.meta_path, "r", encoding="utf-8") as f:
            self.meta = json.load(f)

        dtype = np.dtype(self.meta["dtype"])

        self.data = np.memmap(
            self.bin_path,
            dtype=dtype,
            mode="r",
        )

        if len(self.data) <= block_size + 1:
            raise ValueError("Dataset too small for selected block_size.")

        self.n_samples = (len(self.data) - 1) // block_size

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

    def __getitem__(self, idx: int):
        """
        Input:
            idx: sample index.

        Output:
            x: input tokens.
            y: next-token targets.
        """
        start = idx * self.block_size
        end = start + self.block_size + 1

        chunk = np.asarray(self.data[start:end], dtype=np.int64)

        x = torch.from_numpy(chunk[:-1].copy()).long()
        y = torch.from_numpy(chunk[1:].copy()).long()

        return x, y