# Copyright 2025 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""ExactCoveragePool: exactly-once frame coverage over a bounded episode pool."""

from collections import Counter

import pytest

from lerobot.streaming.episode_pool import ExactCoveragePool

EPISODES = [(0, 5), (1, 3), (2, 8), (3, 1), (4, 6), (5, 4), (6, 7), (7, 2)]
TOTAL = sum(n for _, n in EPISODES)
EXPECTED = Counter((ep, i) for ep, n in EPISODES for i in range(n))


def _drain(pool):
    out, max_resident = [], 0
    while True:
        try:
            out.append(next(pool))
        except StopIteration:
            break
        max_resident = max(max_resident, len(pool.resident))
    return out, max_resident


def test_exact_once_coverage():
    out, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42))
    assert len(out) == TOTAL
    assert Counter(out) == EXPECTED  # every (episode, frame) exactly once, no dups/misses


def test_pool_never_exceeds_size():
    _, max_resident = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42))
    assert max_resident <= 3


def test_deterministic_per_seed_and_epoch():
    a, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7))
    b, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7))
    c, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=8))
    d, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7, epoch=1))
    assert a == b
    assert a != c and a != d  # seed and epoch both change the order
    assert Counter(c) == EXPECTED and Counter(d) == EXPECTED  # ... but coverage is preserved


def test_admission_and_eviction_events():
    pool = ExactCoveragePool(EPISODES, pool_size=3, seed=0)
    admitted_ever, evicted_ever = set(), set()
    # first three episodes admitted at construction
    admitted_ever.update(pool.newly_admitted)
    assert len(admitted_ever) == 3
    while True:
        pool.newly_admitted.clear()
        pool.evicted.clear()
        try:
            next(pool)
        except StopIteration:
            break
        admitted_ever.update(pool.newly_admitted)
        evicted_ever.update(pool.evicted)
    assert admitted_ever == {ep for ep, _ in EPISODES}  # every episode admitted exactly once
    # every episode except the pool_size still resident at the end is evicted on exhaustion
    assert len(evicted_ever) >= len(EPISODES) - 3


def test_uniform_mixing_matches_coupon_collector():
    # 64 equal episodes, pool 64, first 64 draws -> ~64*(1-(1-1/64)^64) ~= 41 distinct
    big = [(e, 100) for e in range(64)]
    pool = ExactCoveragePool(big, pool_size=64, seed=0)
    head = [next(pool)[0] for _ in range(64)]
    assert len(set(head)) >= 30  # far above sequential (=1); ~41 expected


def test_large_epoch_bounded_and_complete():
    big = [(e, 90) for e in range(500)]
    out, max_resident = _drain(ExactCoveragePool(big, pool_size=64, seed=3))
    assert len(out) == 500 * 90
    assert len(set(out)) == 500 * 90  # exactly once
    assert max_resident <= 64


def test_zero_length_episodes_skipped():
    pool = ExactCoveragePool([(0, 3), (1, 0), (2, 2)], pool_size=8, seed=0)
    out, _ = _drain(pool)
    assert Counter(out) == Counter({(0, 0): 1, (0, 1): 1, (0, 2): 1, (2, 0): 1, (2, 1): 1})


def test_byte_aware_admission_never_exceeds_budget():
    sizes = {0: 7, 1: 6, 2: 4, 3: 3, 4: 2}
    pool = ExactCoveragePool(
        EPISODES[:5],
        pool_size=4,
        seed=9,
        episode_byte_sizes=sizes,
        byte_budget=10,
    )
    out = []
    max_resident_bytes = pool.resident_bytes
    while pool.remaining_total:
        out.append(next(pool))
        max_resident_bytes = max(max_resident_bytes, pool.resident_bytes)

    assert Counter(out) == Counter((ep, frame) for ep, count in EPISODES[:5] for frame in range(count))
    assert max_resident_bytes <= 10


def test_byte_aware_admission_rejects_one_oversized_episode():
    with pytest.raises(ValueError, match="Episode 1.*byte budget"):
        ExactCoveragePool(
            [(0, 2), (1, 3)],
            pool_size=2,
            seed=0,
            episode_byte_sizes={0: 4, 1: 11},
            byte_budget=10,
        )


def test_prefetch_candidates_follow_deterministic_pending_frontier():
    pool = ExactCoveragePool(
        EPISODES,
        pool_size=2,
        seed=17,
        episode_byte_sizes={episode: 1 for episode, _ in EPISODES},
        byte_budget=2,
    )

    candidates = pool.prefetch_candidates(3)

    assert len(candidates) == 3
    assert not set(candidates) & set(pool.resident)
    assert candidates == pool.prefetch_candidates(3)


def test_default_sampling_order_is_unchanged() -> None:
    assert list(ExactCoveragePool([(0, 3), (1, 2), (2, 4)], 2, seed=7)) == [
        (0, 0),
        (0, 2),
        (2, 3),
        (2, 2),
        (2, 1),
        (0, 1),
        (1, 1),
        (2, 0),
        (1, 0),
    ]


@pytest.mark.parametrize("seed", [0, 7, 42])
@pytest.mark.parametrize("pool_size", [1, 3, 20])
def test_round_robin_visits_each_round_resident_once(seed: int, pool_size: int) -> None:
    pool = ExactCoveragePool(
        EPISODES,
        pool_size,
        seed=seed,
        sampling_strategy="round_robin",
        episode_byte_sizes=dict(EPISODES),
        byte_budget=12,
    )
    output = []
    admitted = list(pool.newly_admitted)
    evicted = []
    while pool.resident:
        # Admissions during this round must wait until the next round.
        snapshot = set(pool.resident)
        round_episodes = []
        for _ in snapshot:
            pool.newly_admitted.clear()
            pool.evicted.clear()
            sample = next(pool)
            output.append(sample)
            round_episodes.append(sample[0])
            admitted.extend(pool.newly_admitted)
            evicted.extend(pool.evicted)
            assert pool.resident_bytes <= 12
            assert len(pool.resident) <= pool_size
        assert Counter(round_episodes) == Counter(snapshot)
    assert Counter(output) == EXPECTED
    assert Counter(admitted) == Counter(ep for ep, _ in EPISODES)
    assert Counter(evicted) == Counter(admitted)
    with pytest.raises(StopIteration):
        next(pool)


def test_round_robin_replay_seed_and_epoch() -> None:
    def samples(seed: int, epoch: int) -> list[tuple[int, int]]:
        return list(ExactCoveragePool(EPISODES, 3, seed=seed, epoch=epoch, sampling_strategy="round_robin"))

    expected = samples(7, 0)
    assert samples(7, 0) == expected
    assert samples(8, 0) != expected
    assert samples(7, 1) != expected
    restored = ExactCoveragePool(EPISODES, 3, seed=7, sampling_strategy="round_robin")
    for _ in range(13):
        next(restored)
    assert list(restored) == expected[13:]


def test_round_robin_empty_pool() -> None:
    assert list(ExactCoveragePool([(0, 0)], 3, seed=0, sampling_strategy="round_robin")) == []


def test_invalid_sampling_strategy() -> None:
    with pytest.raises(ValueError, match="StreamingSamplingStrategy"):
        ExactCoveragePool(EPISODES, 3, seed=0, sampling_strategy="invalid")
