"""二进制 token 流的批采样。

语料存成一维 token 数组，每次随机截 seq_len+1 的窗口，
前 seq_len 个当输入、后 seq_len 个当标签。省内存、随机性够。
"""

import json
import os
from typing import Tuple

import numpy as np
import torch


class BinDataset:
    def __init__(self, data_dir: str):
        with open(os.path.join(data_dir, "meta.json"), "r", encoding="utf-8") as f:
            self.meta = json.load(f)
        self.dtype = np.dtype(self.meta["dtype"])
        self.data_dir = data_dir
        self._cache = {}

    @property
    def vocab_size(self) -> int:
        return self.meta["vocab_size"]

    def _arr(self, split: str) -> np.ndarray:
        if split not in self._cache:
            path = os.path.join(self.data_dir, f"{split}.bin")
            self._cache[split] = np.memmap(path, dtype=self.dtype, mode="r")
        return self._cache[split]

    def get_batch(self, split: str, batch_size: int, seq_len: int,
                  device: torch.device, generator=None) -> Tuple[torch.Tensor, torch.Tensor]:
        arr = self._arr(split)
        hi = len(arr) - seq_len - 1
        if hi <= 0:
            raise ValueError(f"{split} 语料太短（{len(arr)} tokens），放不下 seq_len={seq_len}")
        ix = torch.randint(hi, (batch_size,), generator=generator)
        x = torch.stack([torch.from_numpy(arr[i:i + seq_len].astype(np.int64)) for i in ix])
        y = torch.stack([torch.from_numpy(arr[i + 1:i + 1 + seq_len].astype(np.int64)) for i in ix])
        if device.type == "cuda":
            return x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
        return x.to(device), y.to(device)