"""Frame-budget dynamic batching: deterministic, bucketed, resumable.

An *epoch* visits every item exactly once. For epoch ``e`` the sampler

1. draws a permutation of all items from ``default_rng([seed, stream, e])``;
2. cuts it into pools of ``pool_size`` items and sorts each pool by length (stable, so ties
   keep the shuffled order) -- items inside a batch therefore have similar length;
3. greedily packs each sorted pool into batches whose padded size
   ``max_len * batch_size`` stays ``<= frame_budget`` (and ``<= max_items`` items);
4. shuffles the order of the batches.

The result depends only on ``(lengths, frame_budget, pool_size, max_items, seed, stream,
epoch)``, so a run can be resumed from a :class:`SamplerPosition` ``(epoch, batch_index)``.
Nothing is ever dropped or split: an item longer than the budget raises
:class:`BudgetError` at construction.

``stream`` separates independent uses (train = 0, eval = 1) that share one ``seed``.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass

import numpy as np

__all__ = ["BatchStream", "BudgetError", "FrameBudgetSampler", "SamplerPosition"]


class BudgetError(ValueError):
    """An item does not fit into the frame budget (it would have to be split or dropped)."""


@dataclass(frozen=True)
class SamplerPosition:
    """Where the stream is: the *next* batch to be produced is ``batch_index`` of ``epoch``."""

    epoch: int = 0
    batch_index: int = 0


class FrameBudgetSampler:
    """Deterministic epoch -> list of batches (each batch: item indices).

    Args:
        lengths: ``num_frames`` of every item, in dataset order (all >= 1).
        frame_budget: maximum padded frames per batch (``max_len * batch_size``).
        seed: non-negative base seed.
        pool_size: items per sorting pool (bigger = more similar lengths, less randomness).
        max_items: optional cap on the batch size.
        stream: independent random stream id (non-negative).

    Raises:
        ValueError: invalid arguments or an empty dataset.
        BudgetError: some item is longer than ``frame_budget``.
    """

    def __init__(
        self,
        lengths: Sequence[int],
        *,
        frame_budget: int,
        seed: int,
        pool_size: int,
        max_items: int | None = None,
        stream: int = 0,
    ) -> None:
        if len(lengths) == 0:
            raise ValueError("cannot sample from an empty dataset")
        if min(lengths) < 1:
            raise ValueError("every length must be >= 1")
        if seed < 0 or stream < 0:
            raise ValueError("seed and stream must be >= 0")
        if frame_budget < 1 or pool_size < 1 or (max_items is not None and max_items < 1):
            raise ValueError("frame_budget, pool_size and max_items must be >= 1")
        too_long = [(i, n) for i, n in enumerate(lengths) if n > frame_budget]
        if too_long:
            index, frames = too_long[0]
            raise BudgetError(
                f"{len(too_long)} item(s) exceed frame_budget={frame_budget}; "
                f"first: item {index} with {frames} frames"
            )
        self._lengths = np.asarray(lengths, dtype=np.int64)
        self.frame_budget = frame_budget
        self.seed = seed
        self.pool_size = pool_size
        self.max_items = max_items
        self.stream = stream
        self._cache: tuple[int, list[list[int]]] | None = None

    @property
    def num_items(self) -> int:
        return int(self._lengths.shape[0])

    def epoch_batches(self, epoch: int) -> list[list[int]]:
        """All batches of ``epoch`` (item indices; each item appears exactly once)."""
        if epoch < 0:
            raise ValueError("epoch must be >= 0")
        if self._cache is not None and self._cache[0] == epoch:
            return self._cache[1]
        rng = np.random.default_rng([self.seed, self.stream, epoch])
        order = rng.permutation(self.num_items)
        batches: list[list[int]] = []
        for start in range(0, self.num_items, self.pool_size):
            pool = order[start : start + self.pool_size]
            pool = pool[np.argsort(self._lengths[pool], kind="stable")]
            current: list[int] = []
            for index in pool.tolist():
                length = int(self._lengths[index])  # ascending, so it is the batch max
                full = bool(current) and (
                    length * (len(current) + 1) > self.frame_budget
                    or (self.max_items is not None and len(current) >= self.max_items)
                )
                if full:
                    batches.append(current)
                    current = []
                current.append(index)
            if current:
                batches.append(current)
        permutation = rng.permutation(len(batches))
        result = [batches[i] for i in permutation.tolist()]
        self._cache = (epoch, result)
        return result

    def num_batches(self, epoch: int) -> int:
        return len(self.epoch_batches(epoch))


class BatchStream:
    """Infinite stream of batches over epochs with a resumable position."""

    def __init__(
        self, sampler: FrameBudgetSampler, position: SamplerPosition | None = None
    ) -> None:
        self._sampler = sampler
        self._epoch = 0
        self._batch_index = 0
        if position is not None:
            self.set_position(position)

    @property
    def position(self) -> SamplerPosition:
        return SamplerPosition(self._epoch, self._batch_index)

    def set_position(self, position: SamplerPosition) -> None:
        """Jump to ``position``.

        Raises:
            ValueError: ``batch_index`` is beyond the number of batches of that epoch.
        """
        count = self._sampler.num_batches(position.epoch)
        if not 0 <= position.batch_index <= count:
            raise ValueError(
                f"batch_index {position.batch_index} outside 0..{count} for epoch {position.epoch}"
            )
        self._epoch, self._batch_index = position.epoch, position.batch_index

    def next_batch(self) -> list[int]:
        """Item indices of the next batch; rolls over to the next epoch when exhausted."""
        batches = self._sampler.epoch_batches(self._epoch)
        if self._batch_index >= len(batches):
            self._epoch += 1
            self._batch_index = 0
            batches = self._sampler.epoch_batches(self._epoch)
        batch = batches[self._batch_index]
        self._batch_index += 1
        return list(batch)