"""Data plumbing for stage-1 training.

* StepDataset : uint16 [N,768] ids + mask_start + len, length-bucketed so each batch
  pads only to the longest member (fixed 768 padding wastes 77% of the compute).
* LMDataset   : random ctx windows out of the packed mathlib token stream.
"""
from __future__ import annotations

import json
import os
import random

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


class StepDataset(Dataset):
    def __init__(self, root: str, split: str, ctx: int = 768, pad_id: int = 0,
                 bucket: int = 512, seed: int = 0):
        self.ids = np.load(f'{root}/{split}_ids.npy', mmap_mode='r')
        self.mask_start = np.load(f'{root}/{split}_mask_start.npy')
        self.lens = np.load(f'{root}/{split}_len.npy')
        self.ctx = ctx
        self.pad_id = pad_id
        self.bucket = bucket
        self.seed = seed
        self.epoch = 0
        self.order = None
        self._reorder()

    def _reorder(self):
        """Sort by length into buckets, shuffle bucket order + inside buckets."""
        rng = random.Random(self.seed + self.epoch)
        idx = np.argsort(self.lens, kind='stable')
        buckets = [idx[i:i + self.bucket] for i in range(0, len(idx), self.bucket)]
        for b in buckets:
            rng.shuffle(b.tolist())
        rng.shuffle(buckets)
        self.order = np.concatenate(buckets) if buckets else idx

    def set_epoch(self, epoch: int):
        self.epoch = epoch
        self._reorder()

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, i):
        j = int(self.order[i])
        n = int(self.lens[j])
        ids = np.asarray(self.ids[j, :n], dtype=np.int64)
        return ids, int(self.mask_start[j])


def collate_steps(batch, pad_id: int = 0, pad_multiple: int = 8):
    lens = [len(b[0]) for b in batch]
    m = max(lens)
    m = min(((m + pad_multiple - 1) // pad_multiple) * pad_multiple, max(lens))
    m = ((m + pad_multiple - 1) // pad_multiple) * pad_multiple
    ids = np.full((len(batch), m), pad_id, dtype=np.int64)
    mask_start = np.zeros(len(batch), dtype=np.int64)
    for k, (seq, ms) in enumerate(batch):
        ids[k, :len(seq)] = seq
        mask_start[k] = min(ms, m)
    return (torch.from_numpy(ids), torch.from_numpy(mask_start))


class LMDataset(Dataset):
    def __init__(self, path: str, ctx: int = 768, length: int = 100_000, seed: int = 0):
        self.tokens = np.load(path, mmap_mode='r')
        self.ctx = ctx
        self.length = length
        self.seed = seed

    def __len__(self):
        return self.length

    def __getitem__(self, i):
        rng = random.Random(self.seed * 1_000_003 + i)
        s = rng.randrange(0, len(self.tokens) - self.ctx - 1)
        seq = np.asarray(self.tokens[s:s + self.ctx + 1], dtype=np.int64)
        return seq[:-1], seq[1:]


def load_specials(tok_dir: str):
    return json.load(open(os.path.join(tok_dir, 'special_tokens.json')))