# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0

"""Deterministic bounded-pool sampling with exact episode-frame coverage."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from enum import Enum

import numpy as np
from numpy.typing import NDArray


class StreamingSamplingStrategy(str, Enum):
    """How the episode pool picks the next anchor among resident episodes."""

    # Pick an episode proportionally to its remaining anchors.
    REMAINING = "remaining"
    # Shuffle resident episodes each round and draw one anchor from each.
    ROUND_ROBIN = "round_robin"


class ExactCoveragePool:
    """Deterministic, exactly-once frame coverage over a byte-cache episode pool.

    A with-replacement pool never guarantees a full
    epoch: frames are drawn randomly and episodes rotate on a fixed cadence. This planner instead
    enumerates *every frame of every episode exactly once per epoch* while keeping at most
    ``pool_size`` episodes resident, so batch mixing stays high but coverage is complete and
    reproducible.

    Mechanics (this is the "evict only when all frames sampled" model):
      - Episodes are admitted in a seeded global permutation until either ``pool_size`` or the
        optional indexed-byte budget is reached.
      - Each resident episode carries a seeded shuffle of its own frame indices.
      - By default each draw picks an episode proportional to its *remaining* frames, uniformly
        sampling the remaining resident frames, not the whole dataset. Alternatively, ``round_robin``
        shuffles the resident episodes each round and pops one frame per episode; newly admitted
        episodes join the next round. This improves episode diversity but is not globally uniform.
      - An episode is evicted only when its last frame is emitted; a new episode is then admitted.
      - The epoch ends when the admission order is exhausted and every resident episode is drained.

    Newly admitted episodes are surfaced via :attr:`newly_admitted` (drain it to drive prefetch)
    and evictions via :attr:`evicted` (drain to release cache bytes). The planner does no I/O and
    is fully unit-testable. It yields ``(episode_index, frame_index)``; the dataset resolves the
    corresponding timestamps from its tabular data.

    Determinism: the order is a pure function of ``(seed, epoch)``, the episode frame counts, and
    sampling strategy and optional byte sizes/budget. Resume is a deterministic fast-forward:
    re-instantiate with the same inputs and skip ``n`` samples (tabular only, no decode).
    """

    def __init__(
        self,
        episode_frame_counts: Sequence[tuple[int, int]],
        pool_size: int,
        *,
        seed: int,
        epoch: int = 0,
        episode_byte_sizes: Mapping[int, int] | None = None,
        byte_budget: int | None = None,
        sampling_strategy: StreamingSamplingStrategy | str = StreamingSamplingStrategy.REMAINING,
    ) -> None:
        """Build a seeded admission plan under optional compressed-byte limits.

        Args:
            episode_frame_counts (`Sequence[tuple[int, int]]`):
                Episode identifiers and their anchor counts. Non-positive counts are omitted.
            pool_size (`int`):
                Maximum resident episodes, clamped to at least one.
            seed (`int`):
                Base seed for episode admission and per-episode anchor permutations.
            epoch (`int`, *optional*, defaults to `0`):
                Coverage epoch combined with the seed to choose the sample order.
            episode_byte_sizes (`Mapping[int, int] | None`, *optional*):
                Indexed compressed-video sizes; required when byte_budget is set.
            byte_budget (`int | None`, *optional*):
                Maximum summed bytes of resident episodes, or None for a slot-only bound.
            sampling_strategy (`StreamingSamplingStrategy | str`, *optional*, defaults to `"remaining"`):
                Remaining-anchor weighting or one anchor per episode in each shuffled round.

        Raises:
            ValueError: If the strategy or byte limits are invalid, or an episode exceeds the budget.
        """
        self.sampling_strategy = StreamingSamplingStrategy(sampling_strategy)
        self._round: list[int] = []
        self._counts = {int(ep): int(n) for ep, n in episode_frame_counts if int(n) > 0}
        self._rng = np.random.default_rng([seed, epoch])
        order = np.array(sorted(self._counts), dtype=np.int64)
        self._rng.shuffle(order)
        self.pool_size = max(1, pool_size)
        self._byte_budget = byte_budget
        if byte_budget is not None and byte_budget <= 0:
            raise ValueError("byte_budget must be positive")
        if byte_budget is not None and episode_byte_sizes is None:
            raise ValueError("episode_byte_sizes are required when byte_budget is set")
        self._byte_sizes = {
            episode: int(episode_byte_sizes[episode]) if episode_byte_sizes is not None else 0
            for episode in self._counts
        }
        if any(size < 0 for size in self._byte_sizes.values()):
            raise ValueError("episode byte sizes must be non-negative")
        if byte_budget is not None:
            oversized = next(
                ((episode, size) for episode, size in self._byte_sizes.items() if size > byte_budget),
                None,
            )
            if oversized is not None:
                episode, size = oversized
                raise ValueError(
                    f"Episode {episode} requires {size} bytes, exceeding the byte budget {byte_budget}"
                )

        self._pending: list[int] = order.tolist()
        self._remaining: dict[int, tuple[NDArray[np.int64], int]] = {}
        self._remaining_total = 0
        self._resident_bytes = 0
        self.newly_admitted: list[int] = []
        self.evicted: list[int] = []
        self._admit_available()

    def _admit_available(self) -> None:
        """Admit pending episodes that fit both the slot and indexed-byte limits."""
        while len(self._remaining) < self.pool_size and self._pending:
            available_bytes = None if self._byte_budget is None else self._byte_budget - self._resident_bytes
            pending_index = next(
                (
                    index
                    for index, episode in enumerate(self._pending)
                    if available_bytes is None or self._byte_sizes[episode] <= available_bytes
                ),
                None,
            )
            if pending_index is None:
                return

            episode = self._pending.pop(pending_index)
            frame_count = self._counts[episode]
            frames: NDArray[np.int64] = np.arange(frame_count, dtype=np.int64)
            self._rng.shuffle(frames)
            self._remaining[episode] = (frames, frame_count)
            self._remaining_total += frame_count
            self._resident_bytes += self._byte_sizes[episode]
            self.newly_admitted.append(episode)

    @property
    def remaining_total(self) -> int:
        """Return the number of unsampled anchors in the currently resident episodes."""
        return self._remaining_total

    @property
    def resident(self) -> list[int]:
        """Return the currently admitted episode indices."""
        return list(self._remaining)

    @property
    def resident_bytes(self) -> int:
        """Return indexed bytes for the currently admitted episodes."""
        return self._resident_bytes

    def prefetch_candidates(self, count: int) -> list[int]:
        """Return the next deterministic pending frontier without admitting it."""
        if count <= 0:
            return []
        return self._pending[:count]

    def __iter__(self) -> ExactCoveragePool:
        """Return this stateful planner as an iterator."""
        return self

    def __next__(self) -> tuple[int, int]:
        """Emit the next episode and frame index in the coverage plan."""
        if self._remaining_total == 0:
            raise StopIteration
        if self.sampling_strategy is StreamingSamplingStrategy.ROUND_ROBIN:
            if not self._round:
                self._round = list(self._remaining)
                self._rng.shuffle(self._round)
            chosen = self._round.pop()
        else:
            # Preserve the default's RNG calls and remaining-frame weighting.
            target = int(self._rng.integers(self._remaining_total))
            chosen = next(iter(self._remaining))
            for ep, (_frames, remaining) in self._remaining.items():
                if target < remaining:
                    chosen = ep
                    break
                target -= remaining
        frames, remaining = self._remaining[chosen]
        remaining -= 1
        frame_index = int(frames[remaining])
        self._remaining_total -= 1
        if remaining == 0:
            del self._remaining[chosen]
            self.evicted.append(chosen)
            self._resident_bytes -= self._byte_sizes[chosen]
            self._admit_available()
        else:
            self._remaining[chosen] = (frames, remaining)
        return chosen, frame_index
