"""Deterministic sampling of training windows from memory-mapped token shards."""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch


def effective_mix(target: dict[str, float], group_tokens: dict[str, int], total_tokens: int,
                  max_epochs: dict[str, float]) -> dict[str, float]:
    """Cap groups at `max_epochs` passes over their data; give the freed share to the other groups."""
    capped = {}
    for g, w in target.items():
        if g in max_epochs:
            cap = max_epochs[g] * group_tokens.get(g, 0) / total_tokens
            if cap < w:
                capped[g] = cap
    free = 1.0 - sum(capped.values())
    rest = {g: w for g, w in target.items() if g not in capped}
    rest_total = sum(rest.values())
    out = dict(capped)
    for g, w in rest.items():
        out[g] = free * w / rest_total
    return {g: out[g] for g in target}


class MixtureLoader:
    """Samples (x, y) windows; batch contents depend only on (seed, step, micro)."""

    def __init__(self, data_dir: Path, split: str, weights: dict[str, float], context_len: int, micro_batch: int,
                 seed: int):
        self.data_dir = Path(data_dir)
        self.manifest = json.loads((self.data_dir / "manifest.json").read_text())
        self.context_len, self.micro_batch, self.seed = context_len, micro_batch, seed
        self.groups: dict[str, tuple[list[np.memmap], np.ndarray]] = {}
        for group, splits in sorted(self.manifest["groups"].items()):
            if split not in splits:
                continue
            arrays, sizes = [], []
            for shard in splits[split]["shards"]:
                if shard["tokens"] < context_len + 1:
                    continue  # too short to hold one window
                # Shard paths are resolved under data_dir so the corpus can be moved (e.g. to Kaggle).
                path = self.data_dir / group / Path(shard["path"]).name
                arrays.append(np.memmap(path, dtype="<u2", mode="r"))
                sizes.append(shard["tokens"] - context_len)
            if arrays:
                self.groups[group] = (arrays, np.asarray(sizes, dtype=np.float64))
        usable = {g: w for g, w in weights.items() if w > 0 and g in self.groups}
        if not usable:
            raise ValueError(f"no usable shards for weights {weights} in split {split!r}")
        total = sum(usable.values())
        self.group_names = sorted(usable)
        self.group_probs = np.asarray([usable[g] / total for g in self.group_names])

    def _window(self, rng: np.random.Generator, group: str) -> np.ndarray:
        arrays, sizes = self.groups[group]
        k = rng.choice(len(arrays), p=sizes / sizes.sum())
        off = int(rng.integers(0, sizes[k]))
        return np.asarray(arrays[k][off : off + self.context_len + 1], dtype=np.int64)

    def _stack(self, rows: list[np.ndarray]) -> tuple[torch.Tensor, torch.Tensor]:
        batch = torch.from_numpy(np.stack(rows))
        return batch[:, :-1].contiguous(), batch[:, 1:].contiguous()

    def get_batch(self, step: int, micro: int) -> tuple[torch.Tensor, torch.Tensor]:
        rng = np.random.default_rng([self.seed, step, micro])
        groups = rng.choice(len(self.group_names), size=self.micro_batch, p=self.group_probs)
        return self._stack([self._window(rng, self.group_names[g]) for g in groups])

    def get_group_batch(self, group: str, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        gid = sorted(self.groups).index(group)
        rng = np.random.default_rng([self.seed, 1_000_000_007, gid, index])
        return self._stack([self._window(rng, group) for _ in range(self.micro_batch)])