"""Deterministic, resumable batch loader over the packed token corpus.

The corpus is one flat uint16 array in which every source already holds exactly its mixture share.
It is cut into 2048-token blocks and visited once in a fixed random order, so a single pass
reproduces the intended mix with no repetition. Step N always sees the same blocks (resume-safe).
"""
from __future__ import annotations

import json
import queue
import threading
from pathlib import Path

import numpy as np
import torch

BLOCK = 2048


class BlockLoader:
    def __init__(self, data_dir: Path, seed: int = 1337, corpus_json: str = "corpus.json", train_file: str = "train.bin"):
        self.dir = Path(data_dir)
        self.meta = json.loads((self.dir / corpus_json).read_text())
        self.mm = np.memmap(self.dir / train_file, dtype=np.uint16, mode="r")
        self.n_blocks = (len(self.mm) - 1) // BLOCK
        self.perm = np.random.default_rng(seed).permutation(self.n_blocks).astype(np.int64)

    def batch(self, step: int, blocks_per_step: int, seq_len: int) -> torch.Tensor:
        """Returns int64 [blocks_per_step * (BLOCK // seq_len), seq_len + 1]: inputs and shifted targets in one array."""
        ids = self.perm[(np.arange(blocks_per_step) + step * blocks_per_step) % self.n_blocks]
        out = np.empty((blocks_per_step, BLOCK // seq_len, seq_len + 1), dtype=np.int64)
        for i, b in enumerate(ids):
            w = self.mm[b * BLOCK : b * BLOCK + BLOCK + 1].astype(np.int64)
            for j in range(BLOCK // seq_len):
                out[i, j] = w[j * seq_len : j * seq_len + seq_len + 1]
        return torch.from_numpy(out.reshape(-1, seq_len + 1))


class StagedLoader:
    """first loader for steps < from_step, then a second corpus from its own beginning (opened when first needed)."""

    def __init__(self, first: BlockLoader, from_step: int, second_dir: Path, seed: int):
        self.first, self.from_step, self.second_dir, self.seed = first, from_step, Path(second_dir), seed
        self.second, self.n_blocks, self.meta = None, first.n_blocks, first.meta

    def ready(self) -> bool:
        return (self.second_dir / "corpus.json").exists()

    def batch(self, step: int, blocks_per_step: int, seq_len: int) -> torch.Tensor:
        if step < self.from_step: return self.first.batch(step, blocks_per_step, seq_len)
        if self.second is None: self.second = BlockLoader(self.second_dir, seed=self.seed)
        return self.second.batch(step - self.from_step, blocks_per_step, seq_len)


class Prefetcher:
    """Background thread that prepares upcoming steps' batches."""

    def __init__(self, loader: BlockLoader, start_step: int, blocks_per_step: int, seq_len_fn, depth: int = 3):
        self.q: queue.Queue = queue.Queue(maxsize=depth)
        self.loader, self.bps, self.seq_len_fn = loader, blocks_per_step, seq_len_fn
        self.step = start_step
        self.stop = False
        self.t = threading.Thread(target=self._run, daemon=True)
        self.t.start()

    def _run(self):
        s = self.step
        while not self.stop:
            T = self.seq_len_fn(s)
            b = self.loader.batch(s, self.bps, T).pin_memory()
            self.q.put((s, T, b))
            s += 1

    def get(self):
        return self.q.get()


class ValSet:
    """Fixed held-out token windows, overall and per source."""

    def __init__(self, data_dir: Path, seq_len: int = 2048, per_source_windows: int = 64):
        meta = json.loads((Path(data_dir) / "corpus.json").read_text())
        mm = np.memmap(Path(data_dir) / "val.bin", dtype=np.uint16, mode="r")
        self.sets: dict[str, torch.Tensor] = {}
        for name, s in meta["sources"].items():
            n = (s["val_tokens"] - 1) // seq_len
            if n <= 0:
                continue
            take = min(n, per_source_windows)
            starts = np.linspace(0, n - 1, take).astype(int)
            arr = np.stack([mm[s["val_start"] + i * seq_len : s["val_start"] + i * seq_len + seq_len + 1].astype(np.int64) for i in starts])
            self.sets[name] = torch.from_numpy(arr)