tests/test_frame_sampling.py

"""Frame sampling and video decoding: contract tests against the reference pipeline.

The model and the spatial transform are verified elsewhere. What is left is the
part that is *your* code: which frame indices a clip is built from, and what the
decoder returns for them. Both fail silently — the model happily consumes wrong
frames in the wrong colour order — so they get their own ground truth here.

The tests come in two groups.

**Self-contained.** They synthesise a video in which frame `i` is a solid colour
encoding `i`, so a decoded frame states its own index. Nothing needs to be
plugged in for these to run.

**Contract.** They compare your sampler and decoder against the reference. Point
the environment at them as `module:function`:

    VJEPA21_USER_SAMPLER=mypkg.data:clip_indices \\
    VJEPA21_USER_DECODER=mypkg.data:decode_frames \\
    python -m pytest test_frame_sampling.py -s -q

Expected signatures:

    clip_indices(video_len, frames_per_clip, frame_step, num_clips=1) -> Sequence[Sequence[int]]
    decode_frames(path, indices) -> np.ndarray  # (T, H, W, 3), uint8, RGB
"""

from __future__ import annotations

import importlib
import math
import os
import subprocess
import tempfile

import numpy as np
import pytest

USER_SAMPLER = os.environ.get("VJEPA21_USER_SAMPLER", "")
USER_DECODER = os.environ.get("VJEPA21_USER_DECODER", "")

N_FRAMES = 120
FRAME_H, FRAME_W = 64, 96


# --- reference index arithmetic ---------------------------------------------


def official_clip_indices(
    video_len: int,
    frames_per_clip: int,
    frame_step: int,
    num_clips: int = 1,
    allow_clip_overlap: bool = False,
    random_clip_sampling: bool = False,
) -> list[np.ndarray]:
    """Transcription of `VideoDataset.loadvideo_decord` index selection.

    Source: `src/datasets/video_dataset.py`, the block after `vr.seek(0)`.
    Deterministic when `random_clip_sampling=False`, which is the evaluation
    setting.
    """
    fpc, fstp = frames_per_clip, frame_step
    clip_len = int(fpc * fstp)
    partition_len = video_len // num_clips

    clip_indices = []
    for i in range(num_clips):
        if partition_len > clip_len:
            end_indx = clip_len
            if random_clip_sampling:
                end_indx = np.random.randint(clip_len, partition_len)
            start_indx = end_indx - clip_len
            indices = np.linspace(start_indx, end_indx, num=fpc)
            indices = np.clip(indices, start_indx, end_indx - 1).astype(np.int64)
            indices = indices + i * partition_len
        elif not allow_clip_overlap:
            indices = np.linspace(0, partition_len, num=partition_len // fstp)
            indices = np.concatenate(
                (indices, np.ones(fpc - partition_len // fstp) * partition_len)
            )
            indices = np.clip(indices, 0, partition_len - 1).astype(np.int64)
            indices = indices + i * partition_len
        else:
            sample_len = min(clip_len, video_len) - 1
            indices = np.linspace(0, sample_len, num=sample_len // fstp)
            indices = np.concatenate(
                (indices, np.ones(fpc - sample_len // fstp) * sample_len)
            )
            indices = np.clip(indices, 0, sample_len - 1).astype(np.int64)
            clip_step = 0
            if video_len > clip_len:
                clip_step = (video_len - clip_len) // (num_clips - 1)
            indices = indices + i * clip_step
        clip_indices.append(indices)
    return clip_indices


def official_frame_step_from_fps(video_fps: float, target_fps: int) -> int:
    """`fstp = math.ceil(avg_fps) // target_fps` — note the ceil, then floor div."""
    return math.ceil(video_fps) // target_fps


# --- synthetic ground-truth video -------------------------------------------


def _colour_for(index: int) -> tuple[int, int, int]:
    """Frame `index` is a solid colour that encodes it. R alone identifies the
    frame; G and B are set so a red/blue swap cannot go unnoticed."""
    return (index * 2 % 256, 40, 210)


@pytest.fixture(scope="module")
def indexed_video():
    """A losslessly encoded video whose frames state their own index."""
    if not _have("ffmpeg"):
        pytest.skip("ffmpeg is required to synthesise the reference video")

    tmpdir = tempfile.mkdtemp()
    path = os.path.join(tmpdir, "indexed.mkv")
    raw = np.zeros((N_FRAMES, FRAME_H, FRAME_W, 3), dtype=np.uint8)
    for i in range(N_FRAMES):
        raw[i, :, :] = _colour_for(i)

    subprocess.run(
        [
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
            "-f", "rawvideo", "-pix_fmt", "rgb24",
            "-s", f"{FRAME_W}x{FRAME_H}", "-r", "30", "-i", "pipe:0",
            "-c:v", "ffv1", "-pix_fmt", "gbrp", path,
        ],
        input=raw.tobytes(), check=True,
    )
    yield path, raw


def _have(binary: str) -> bool:
    from shutil import which

    return which(binary) is not None


def _load(spec: str):
    """Resolve a `module:function` spec, with a readable error when it is wrong."""
    module_name, _, attribute = spec.partition(":")
    if not attribute:
        pytest.fail(f"{spec!r} is not in `module:function` form, e.g. `video_io:clip_indices`")
    try:
        module = importlib.import_module(module_name)
    except ImportError as exc:
        pytest.fail(
            f"cannot import {module_name!r} from {spec!r}: {exc}.\n"
            "This must point at your own code. If you have not written a dataloader yet, "
            "use the reference implementation shipped alongside these tests:\n"
            "    VJEPA21_USER_SAMPLER=video_io:clip_indices "
            "VJEPA21_USER_DECODER=video_io:decode_frames\n"
            "run from the directory containing video_io.py, or with it on PYTHONPATH."
        )
    if not hasattr(module, attribute):
        pytest.fail(f"{module_name!r} has no attribute {attribute!r}")
    return getattr(module, attribute)


def _decord_available() -> bool:
    try:
        import decord  # noqa: F401
    except ImportError:
        return False
    return True


def _require_decord():
    if not _decord_available():
        pytest.skip("decord is not installed; `pip install decord` to enable this check")


def decode_with_decord(path: str, indices) -> np.ndarray:
    _require_decord()
    from decord import VideoReader, cpu

    reader = VideoReader(path, num_threads=-1, ctx=cpu(0))
    reader.seek(0)
    return reader.get_batch(list(indices)).asnumpy()


def _decode_any(path: str, indices) -> np.ndarray:
    """Decode with whatever backend is present, for the ground-truth check.

    Availability is probed by importing rather than by calling
    `decode_with_decord`: `pytest.skip` raises a `BaseException` subclass, so a
    `try/except Exception` around it would let the skip escape and quietly
    disable this check on machines without decord — which is exactly the kind of
    silent no-op the rest of this file exists to catch.
    """
    if _decord_available():
        return decode_with_decord(path, indices)
    import cv2

    wanted, frames, capture, position = set(int(i) for i in indices), {}, cv2.VideoCapture(path), 0
    while len(frames) < len(wanted):
        ok, frame = capture.read()
        if not ok:
            break
        if position in wanted:
            frames[position] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        position += 1
    capture.release()
    return np.stack([frames[int(i)] for i in indices])


def index_of(frame: np.ndarray) -> int:
    """Recover the frame index from a decoded frame, via the red channel."""
    return int(round(float(np.median(frame[..., 0])) / 2))


# --- 1. the arithmetic ------------------------------------------------------


def test_official_sampling_is_not_a_strided_range():
    """The reference spreads `frames_per_clip` samples across `fpc * frame_step`
    with `linspace`, so the effective stride is `fpc*fstp/(fpc-1)`, not `fstp`.

    Reimplementing it as `range(0, fpc*fstp, fstp)` is the intuitive reading and
    it is wrong: at 16 frames with step 4 it picks a different frame 12 times out
    of 16. The clip still looks plausible, which is what makes it dangerous.
    """
    fpc, fstp = 16, 4
    official = official_clip_indices(300, fpc, fstp)[0]
    naive = np.arange(0, fpc * fstp, fstp)

    differing = int((official != naive).sum())
    print(f"\n[sampling] official {official.tolist()}")
    print(f"[sampling] naive    {naive.tolist()}")
    print(f"[sampling] differing frames: {differing}/{fpc}")
    assert differing == 12
    assert official.max() == fpc * fstp - 1


@pytest.mark.parametrize("num_clips", [1, 2, 3])
def test_clips_are_disjoint_and_in_range(num_clips):
    video_len, fpc, fstp = 300, 16, 4
    clips = official_clip_indices(video_len, fpc, fstp, num_clips=num_clips)
    assert len(clips) == num_clips
    for clip in clips:
        assert len(clip) == fpc
        assert clip.min() >= 0 and clip.max() < video_len
        assert (np.diff(clip) >= 0).all(), "indices must be non-decreasing"
    starts = [int(c[0]) for c in clips]
    assert starts == sorted(starts)


def test_short_video_pads_with_the_last_frame():
    """When a partition is shorter than a clip the reference repeats its final
    frame rather than wrapping around or raising."""
    clip = official_clip_indices(video_len=40, frames_per_clip=16, frame_step=4)[0]
    assert len(clip) == 16
    assert clip.max() <= 39
    assert (clip[-1] == clip[-2]) or (clip == clip.max()).sum() > 1


@pytest.mark.parametrize(
    "video_fps,target_fps,expected", [(30.0, 4, 7), (29.97, 4, 7), (25.0, 4, 6), (60.0, 4, 15)]
)
def test_frame_step_from_fps(video_fps, target_fps, expected):
    """`math.ceil` on the average fps before the floor division: 29.97 fps
    behaves like 30, not like 29."""
    assert official_frame_step_from_fps(video_fps, target_fps) == expected


# --- 2. the decoder ---------------------------------------------------------


def test_synthetic_video_is_recoverable(indexed_video):
    """Sanity check on the ground truth itself before it is used to judge anyone."""
    path, raw = indexed_video
    frames = _decode_any(path, range(N_FRAMES))
    assert frames.shape == raw.shape and frames.dtype == np.uint8
    recovered = [index_of(f) for f in frames]
    assert recovered == list(range(N_FRAMES))


@pytest.mark.skipif(not USER_DECODER, reason="set VJEPA21_USER_DECODER=module:function")
def test_user_decoder_returns_rgb(indexed_video):
    """A BGR decoder — `cv2.VideoCapture` returns BGR — feeds the model channel
    swapped. Nothing errors; the features are simply wrong."""
    path, _ = indexed_video
    frames = np.asarray(_load(USER_DECODER)(path, [0]))
    red, _green, blue = frames[0].reshape(-1, 3).mean(axis=0)
    print(f"\n[decoder] frame 0 mean RGB = ({red:.0f}, {_green:.0f}, {blue:.0f}); "
          f"expected ≈ {_colour_for(0)}")
    assert blue > red, "channels look swapped: this is BGR, the model expects RGB"


@pytest.mark.skipif(not USER_DECODER, reason="set VJEPA21_USER_DECODER=module:function")
def test_user_decoder_output_contract(indexed_video):
    path, _ = indexed_video
    frames = np.asarray(_load(USER_DECODER)(path, [0, 5, 10]))
    assert frames.shape == (3, FRAME_H, FRAME_W, 3), f"expected (T, H, W, 3), got {frames.shape}"
    assert frames.dtype == np.uint8, f"expected uint8, got {frames.dtype}"


@pytest.mark.skipif(not USER_DECODER, reason="set VJEPA21_USER_DECODER=module:function")
def test_user_decoder_returns_the_requested_frames(indexed_video):
    """Off-by-one seeking, keyframe snapping and dropped frames all land here."""
    path, _ = indexed_video
    wanted = [0, 1, 17, 42, 63, 99, N_FRAMES - 1]
    frames = np.asarray(_load(USER_DECODER)(path, wanted))
    got = [index_of(f) for f in frames]
    print(f"\n[decoder] requested {wanted}\n[decoder] received  {got}")
    assert got == wanted


@pytest.mark.skipif(not USER_DECODER, reason="set VJEPA21_USER_DECODER=module:function")
def test_user_decoder_matches_decord(indexed_video):
    """Pixel-level agreement with the decoder the reference pipeline uses."""
    path, _ = indexed_video
    wanted = [0, 7, 31, 64, 111]
    mine = np.asarray(_load(USER_DECODER)(path, wanted)).astype(np.int16)
    theirs = decode_with_decord(path, wanted).astype(np.int16)
    diff = np.abs(mine - theirs)
    print(f"\n[decoder] max|Δ| vs decord = {diff.max()}/255   mean = {diff.mean():.4f}/255")
    assert diff.max() <= 2, "decoders disagree beyond codec rounding"


@pytest.mark.skipif(not USER_DECODER, reason="set VJEPA21_USER_DECODER=module:function")
def test_user_decoder_is_deterministic(indexed_video):
    path, _ = indexed_video
    decode = _load(USER_DECODER)
    first = np.asarray(decode(path, [3, 14, 15, 92]))
    second = np.asarray(decode(path, [3, 14, 15, 92]))
    assert np.array_equal(first, second)


# --- 3. the sampler ---------------------------------------------------------


@pytest.mark.skipif(not USER_SAMPLER, reason="set VJEPA21_USER_SAMPLER=module:function")
@pytest.mark.parametrize("video_len,fpc,fstp,num_clips", [
    (300, 16, 4, 1), (300, 16, 4, 3), (120, 16, 4, 1), (40, 16, 4, 1), (1000, 32, 2, 2),
])
def test_user_sampler_matches_official(video_len, fpc, fstp, num_clips):
    expected = official_clip_indices(video_len, fpc, fstp, num_clips=num_clips)
    got = _load(USER_SAMPLER)(video_len, fpc, fstp, num_clips=num_clips)
    got = [np.asarray(clip, dtype=np.int64) for clip in got]

    assert len(got) == len(expected), f"expected {len(expected)} clips, got {len(got)}"
    for i, (mine, reference) in enumerate(zip(got, expected)):
        if not np.array_equal(mine, reference):
            print(f"\n[sampler] clip {i} expected {reference.tolist()}")
            print(f"[sampler] clip {i} got      {mine.tolist()}")
        assert np.array_equal(mine, reference), f"clip {i} differs"


@pytest.mark.skipif(
    not (USER_SAMPLER and USER_DECODER), reason="set both VJEPA21_USER_* variables"
)
def test_user_pipeline_end_to_end(indexed_video):
    """Sampler and decoder together must land on the frames the reference would."""
    path, _ = indexed_video
    fpc, fstp = 16, 4
    expected = official_clip_indices(N_FRAMES, fpc, fstp)[0]
    clips = _load(USER_SAMPLER)(N_FRAMES, fpc, fstp, num_clips=1)
    frames = np.asarray(_load(USER_DECODER)(path, list(clips[0])))
    got = [index_of(f) for f in frames]
    print(f"\n[pipeline] expected {expected.tolist()}\n[pipeline] got      {got}")
    assert got == expected.tolist()


# --- 4. multi-clip: coverage and aggregation --------------------------------

import sys as _sys  # noqa: E402

_sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def _video_io():
    try:
        import video_io
    except ImportError:
        pytest.skip("video_io.py not importable; run from the repository root")
    return video_io


def test_partitioned_sampling_leaves_long_videos_mostly_unseen():
    """The reference sampler was designed for short single-label clips.

    On a 10-second video eight segments see almost everything. On a two-minute
    surveillance video the same settings see 14% of it, and leave a blind gap of
    386 frames — thirteen seconds during which an event is not observed at all.
    """
    io = _video_io()
    short = io.temporal_coverage(io.clip_indices(300, 16, 4, num_clips=8), 300)
    long = io.temporal_coverage(io.clip_indices(3600, 16, 4, num_clips=8), 3600)

    print(f"\n[coverage] 10 s video : {short['covered_fraction']*100:.1f}% seen, "
          f"max gap {short['max_gap']} frames")
    print(f"[coverage] 2 min video: {long['covered_fraction']*100:.1f}% seen, "
          f"max gap {long['max_gap']} frames")

    assert short["covered_fraction"] > 0.95
    assert long["covered_fraction"] < 0.20
    assert long["max_gap"] > 300


@pytest.mark.parametrize("video_len", [300, 3600, 9000])
def test_dense_grid_covers_everything(video_len):
    """A sliding grid leaves no gap, which is what frame-level scoring needs."""
    io = _video_io()
    clips = io.dense_clip_indices(video_len, 16, 4)
    metrics = io.temporal_coverage(clips, video_len)
    print(f"\n[dense] {video_len} frames -> {len(clips)} clips, "
          f"{metrics['covered_fraction']*100:.1f}% covered")
    assert metrics["covered_fraction"] == 1.0
    assert metrics["max_gap"] == 0
    assert all(c.max() < video_len for c in clips)


def test_dense_grid_stride_controls_overlap():
    io = _video_io()
    contiguous = io.dense_clip_indices(3600, 16, 4)
    overlapping = io.dense_clip_indices(3600, 16, 4, stride=32)
    assert len(overlapping) > len(contiguous)
    assert io.temporal_coverage(overlapping, 3600)["covered_fraction"] == 1.0


def test_aggregation_averages_probabilities_not_logits():
    """The reference averages softmax outputs. Averaging logits is a different
    estimator and can rank classes differently."""
    io = _video_io()
    views = [np.array([[6.0, 0.0, 0.0]]), np.array([[0.0, 2.0, 2.4]])]

    probabilities = io.aggregate_predictions(views)
    assert np.allclose(probabilities.sum(axis=-1), 1.0)

    logit_mean = np.mean(views, axis=0)[0]
    logit_mean = np.exp(logit_mean - logit_mean.max())
    logit_mean /= logit_mean.sum()

    print(f"\n[aggregate] probability mean {np.round(probabilities[0], 4)}")
    print(f"[aggregate] logit mean       {np.round(logit_mean, 4)}")
    assert not np.allclose(probabilities[0], logit_mean, atol=1e-3)


def test_aggregation_is_order_independent():
    io = _video_io()
    views = [np.random.randn(2, 5) for _ in range(4)]
    a = io.aggregate_predictions(views)
    b = io.aggregate_predictions(views[::-1])
    assert np.allclose(a, b)


@pytest.mark.parametrize("reduce", ["max", "mean", "first"])
def test_clip_scores_reach_every_frame(reduce):
    """Frame-level AUC and AP need a score for every frame, including those no
    clip covered."""
    io = _video_io()
    video_len = 3600
    clips = io.dense_clip_indices(video_len, 16, 4)
    scores = np.linspace(0, 1, len(clips))

    frame_scores = io.clip_scores_to_frame_scores(clips, scores, video_len, reduce=reduce)
    assert frame_scores.shape == (video_len,)
    assert np.isfinite(frame_scores).all()
    assert frame_scores.min() >= scores.min() - 1e-9
    assert frame_scores.max() <= scores.max() + 1e-9


def test_max_reduction_propagates_a_single_high_clip():
    io = _video_io()
    video_len = 1000
    clips = io.dense_clip_indices(video_len, 16, 4)
    scores = np.zeros(len(clips))
    scores[3] = 1.0
    frame_scores = io.clip_scores_to_frame_scores(clips, scores, video_len, reduce="max")
    flagged = int((frame_scores > 0.5).sum())
    print(f"\n[scores] one clip at 1.0 flags {flagged} frames")
    assert flagged >= 64
    assert flagged < video_len