"""Video I/O compatible with the V-JEPA 2.1 reference pipeline.

Two functions, both transcribed from `src/datasets/video_dataset.py` in
`facebookresearch/vjepa2`:

    clip_indices(...)   which frames a clip is built from
    decode_frames(...)  what those frame indices decode to, as RGB uint8

Drop this into your project, or use it as the specification for your own
implementation. `tests/test_frame_sampling.py` checks either against the
reference:

    VJEPA21_USER_SAMPLER=video_io:clip_indices \\
    VJEPA21_USER_DECODER=video_io:decode_frames \\
    python -m pytest tests/test_frame_sampling.py -s -q
"""

from __future__ import annotations

import math
from typing import Sequence

import numpy as np


# --- sampling ---------------------------------------------------------------


def 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,
    rng: np.random.Generator | None = None,
) -> list[np.ndarray]:
    """Frame indices for `num_clips` clips, as the reference pipeline picks them.

    The video is split into `num_clips` equal partitions and one clip is taken
    from each. Note that the frames within a clip are placed with `np.linspace`
    over a window of `frames_per_clip * frame_step`, so the effective stride is
    `fpc*fstp/(fpc-1)` rather than `frame_step`: at 16 frames with step 4 this
    differs from `range(0, 64, 4)` on 12 of the 16 indices.

    Args:
        video_len: Total number of frames in the video.
        frames_per_clip: Frames to return per clip.
        frame_step: Nominal stride; sets the window length, see above.
        num_clips: Number of clips, one per partition.
        allow_clip_overlap: When a partition is shorter than a clip, allow
            successive clips to overlap instead of padding within the partition.
        random_clip_sampling: Random window inside the partition. Leave False
            for evaluation and for a reproducible feature store.
        rng: Optional generator, used only when `random_clip_sampling` is True.

    Returns:
        One int64 array of length `frames_per_clip` per clip.
    """
    if video_len <= 0:
        raise ValueError(f"video_len must be positive, got {video_len}")
    if frames_per_clip <= 0 or frame_step <= 0:
        raise ValueError("frames_per_clip and frame_step must be positive")

    fpc, fstp = frames_per_clip, frame_step
    clip_len = int(fpc * fstp)
    partition_len = video_len // num_clips
    draw = rng if rng is not None else np.random

    out = []
    for i in range(num_clips):
        if partition_len > clip_len:
            end = clip_len
            if random_clip_sampling:
                end = int(draw.integers(clip_len, partition_len)) if rng is not None \
                    else int(np.random.randint(clip_len, partition_len))
            start = end - clip_len
            idx = np.linspace(start, end, num=fpc)
            idx = np.clip(idx, start, end - 1).astype(np.int64) + i * partition_len
        elif not allow_clip_overlap:
            idx = np.linspace(0, partition_len, num=partition_len // fstp)
            idx = np.concatenate((idx, np.ones(fpc - partition_len // fstp) * partition_len))
            idx = np.clip(idx, 0, partition_len - 1).astype(np.int64) + i * partition_len
        else:
            sample_len = min(clip_len, video_len) - 1
            idx = np.linspace(0, sample_len, num=sample_len // fstp)
            idx = np.concatenate((idx, np.ones(fpc - sample_len // fstp) * sample_len))
            idx = np.clip(idx, 0, sample_len - 1).astype(np.int64)
            step = (video_len - clip_len) // (num_clips - 1) if video_len > clip_len else 0
            idx = idx + i * step
        out.append(idx)
    return out


def frame_step_for_fps(video_fps: float, target_fps: int) -> int:
    """Stride that samples a video at `target_fps`.

    The reference rounds the source rate *up* before the integer division, so
    29.97 fps behaves like 30 rather than like 29.
    """
    step = math.ceil(video_fps) // target_fps
    if step < 1:
        raise ValueError(
            f"target_fps={target_fps} exceeds the source rate {video_fps}; step would be 0"
        )
    return step


# --- decoding ---------------------------------------------------------------


def available_backends() -> list[str]:
    backends = []
    for name, module in (("decord", "decord"), ("pyav", "av"), ("opencv", "cv2")):
        try:
            __import__(module)
            backends.append(name)
        except ImportError:
            continue
    return backends


def decode_frames(path: str, indices: Sequence[int], backend: str = "auto") -> np.ndarray:
    """Decode the given frame indices as RGB.

    Args:
        path: Video file.
        indices: Frame indices, as returned by `clip_indices`. Repeats are
            allowed and preserved, which matters because the reference pads
            short clips by repeating the last frame.
        backend: "decord", "pyav", "opencv", or "auto".

    Returns:
        `(len(indices), H, W, 3)` uint8 in **RGB** order.

    The colour order is the point of this function. `cv2.VideoCapture` hands
    back BGR; feeding that to the model swaps red and blue, nothing raises, and
    the features are quietly wrong.
    """
    wanted = [int(i) for i in indices]
    if not wanted:
        raise ValueError("indices is empty")
    if min(wanted) < 0:
        raise ValueError(f"negative frame index: {min(wanted)}")

    if backend == "auto":
        found = available_backends()
        if not found:
            raise RuntimeError(
                "no video backend available; install one of: "
                "decord, av (PyAV), opencv-python"
            )
        backend = found[0]

    if backend == "decord":
        return _decode_decord(path, wanted)
    if backend == "pyav":
        return _decode_pyav(path, wanted)
    if backend == "opencv":
        return _decode_opencv(path, wanted)
    raise ValueError(f"unknown backend {backend!r}")


def _decode_decord(path: str, wanted: list[int]) -> np.ndarray:
    from decord import VideoReader, cpu

    reader = VideoReader(path, num_threads=-1, ctx=cpu(0))
    reader.seek(0)  # the reference rewinds before sampling
    return reader.get_batch(wanted).asnumpy()


def _decode_pyav(path: str, wanted: list[int]) -> np.ndarray:
    import av

    needed = set(wanted)
    frames: dict[int, np.ndarray] = {}
    with av.open(path) as container:
        stream = container.streams.video[0]
        stream.thread_type = "AUTO"
        for position, frame in enumerate(container.decode(stream)):
            if position in needed:
                frames[position] = frame.to_ndarray(format="rgb24")
                if len(frames) == len(needed):
                    break
    _check_complete(frames, wanted, path)
    return np.stack([frames[i] for i in wanted])


def _decode_opencv(path: str, wanted: list[int]) -> np.ndarray:
    import cv2

    needed = set(wanted)
    frames: dict[int, np.ndarray] = {}
    capture = cv2.VideoCapture(path)
    try:
        position = 0
        # Sequential decoding: CAP_PROP_POS_FRAMES snaps to keyframes on several
        # codecs, which silently returns a neighbouring frame.
        while len(frames) < len(needed):
            ok, frame = capture.read()
            if not ok:
                break
            if position in needed:
                frames[position] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            position += 1
    finally:
        capture.release()
    _check_complete(frames, wanted, path)
    return np.stack([frames[i] for i in wanted])


def _check_complete(frames: dict, wanted: list[int], path: str) -> None:
    missing = sorted(set(wanted) - set(frames))
    if missing:
        raise RuntimeError(
            f"{path}: could not decode frames {missing[:8]}"
            f"{'...' if len(missing) > 8 else ''}; the video has {len(frames)} "
            "decodable frames at or before those positions"
        )


__all__ = [
    "clip_indices",
    "dense_clip_indices",
    "frame_step_for_fps",
    "decode_frames",
    "available_backends",
    "aggregate_predictions",
    "temporal_coverage",
    "clip_scores_to_frame_scores",
]


# --- multi-clip: aggregation, coverage, dense grids --------------------------


def aggregate_predictions(logits_per_view: Sequence[np.ndarray]) -> np.ndarray:
    """Combine the predictions of several clips or spatial views of one video.

    Matches the reference evaluation loop, which averages **softmax
    probabilities**, not logits:
    `sum(F.softmax(o, dim=1) for o in outputs) / len(outputs)`
    (`evals/video_classification_frozen/eval.py`).

    Averaging logits instead is a different estimator — it is a geometric rather
    than an arithmetic mean over probabilities — and it lets one confident view
    dominate. The two agree only when the views agree.

    Args:
        logits_per_view: One `(batch, num_classes)` array per clip or view.

    Returns:
        `(batch, num_classes)` probabilities that sum to one.
    """
    if not logits_per_view:
        raise ValueError("logits_per_view is empty")
    probabilities = []
    for logits in logits_per_view:
        logits = np.asarray(logits, dtype=np.float64)
        shifted = logits - logits.max(axis=-1, keepdims=True)
        exponentiated = np.exp(shifted)
        probabilities.append(exponentiated / exponentiated.sum(axis=-1, keepdims=True))
    return np.mean(probabilities, axis=0)


def temporal_coverage(clips: Sequence[np.ndarray], video_len: int) -> dict:
    """How much of a video a set of clips actually looks at.

    The partitioned sampling used by the reference evaluation was designed for
    short single-label videos. On a long one it leaves most of the timeline
    unseen, and an event shorter than `max_gap` can fall entirely between two
    clips without any of it being observed.

    Returns:
        `covered_fraction`, `max_gap` and `num_gaps`, in frames.
    """
    seen = np.zeros(video_len, dtype=bool)
    for clip in clips:
        clip = np.asarray(clip)
        seen[int(clip.min()) : int(clip.max()) + 1] = True

    gaps, run = [], 0
    for visible in seen:
        if visible:
            if run:
                gaps.append(run)
            run = 0
        else:
            run += 1
    if run:
        gaps.append(run)

    return {
        "covered_fraction": float(seen.mean()),
        "max_gap": int(max(gaps)) if gaps else 0,
        "num_gaps": len(gaps),
    }


def dense_clip_indices(
    video_len: int,
    frames_per_clip: int,
    frame_step: int,
    stride: int | None = None,
) -> list[np.ndarray]:
    """A sliding grid of clips covering the whole video.

    Use this instead of `clip_indices` when you need a score per position in
    time rather than one prediction per video — temporal anomaly detection,
    action localisation, anything scored frame by frame. `clip_indices`
    partitions the video and samples one clip per partition, which is the right
    thing for classifying a short video and the wrong thing here.

    Args:
        stride: Frames between the start of consecutive clips. Defaults to the
            clip window `frames_per_clip * frame_step`, giving contiguous
            non-overlapping clips. A smaller value overlaps them, which raises
            temporal resolution at proportional cost.

    Returns:
        Clips in temporal order, the last one clamped to the end of the video.
    """
    window = frames_per_clip * frame_step
    stride = window if stride is None else stride
    if stride <= 0:
        raise ValueError(f"stride must be positive, got {stride}")

    starts = list(range(0, max(video_len - window, 0) + 1, stride))
    if not starts:
        starts = [0]
    if starts[-1] + window < video_len:
        starts.append(video_len - window)

    clips = []
    for start in starts:
        idx = np.linspace(start, start + window, num=frames_per_clip)
        clips.append(np.clip(idx, 0, video_len - 1).astype(np.int64))
    return clips


def clip_scores_to_frame_scores(
    clips: Sequence[np.ndarray],
    scores: Sequence[float],
    video_len: int,
    reduce: str = "max",
) -> np.ndarray:
    """Spread clip-level scores back over frames, for frame-level metrics.

    Frame-level AUC on UCF-Crime and average precision on XD-Violence are
    computed per frame, so a clip score has to be assigned to the frames the
    clip covers. Overlapping clips give a frame several scores; `reduce` picks
    between them. Frames covered by no clip keep the score of the nearest
    covered frame, so the output is dense.

    Args:
        reduce: "max" (an anomaly anywhere in the window marks the window),
            "mean" (smoother, blunter) or "first".
    """
    if len(clips) != len(scores):
        raise ValueError(f"{len(clips)} clips but {len(scores)} scores")

    accumulated = np.zeros(video_len, dtype=np.float64)
    counts = np.zeros(video_len, dtype=np.int64)
    assigned = np.zeros(video_len, dtype=bool)

    for clip, score in zip(clips, scores):
        clip = np.asarray(clip)
        lo, hi = int(clip.min()), int(clip.max()) + 1
        if reduce == "max":
            accumulated[lo:hi] = np.where(
                assigned[lo:hi], np.maximum(accumulated[lo:hi], score), score
            )
        elif reduce == "mean":
            accumulated[lo:hi] += score
        elif reduce == "first":
            accumulated[lo:hi] = np.where(assigned[lo:hi], accumulated[lo:hi], score)
        else:
            raise ValueError(f"unknown reduce {reduce!r}")
        counts[lo:hi] += 1
        assigned[lo:hi] = True

    if reduce == "mean":
        accumulated[counts > 0] /= counts[counts > 0]

    if not assigned.all():
        covered = np.flatnonzero(assigned)
        if covered.size == 0:
            raise ValueError("no frame was covered by any clip")
        nearest = covered[np.abs(np.subtract.outer(np.arange(video_len), covered)).argmin(axis=1)]
        accumulated = np.where(assigned, accumulated, accumulated[nearest])
    return accumulated