"""Fast video processor for V-JEPA 2.1.

Reproduces the evaluation-time transform of the reference implementation
(`evals/video_classification_frozen/utils.py::EvalVideoTransform`):

    resize the short side to `crop_size` -> square crop of `crop_size`
    -> scale to [0, 1] -> normalise with ImageNet statistics

Two things are worth stating precisely.

*Which reference transform this is.* `make_transforms` picks `EvalVideoTransform`
when `num_views_per_clip > 1` and `VideoTransform(training=False)` otherwise, and
only the first resizes the short side to exactly `crop_size` — the second uses
`crop_size * 256 / 224`, like the stock `VJEPA2VideoProcessor`. Every V-JEPA 2.1
frozen-probe config in `configs/eval_2_1/` sets `num_views_per_segment: 3`, so
`EvalVideoTransform` is the transform the 2.1 evaluations actually use and
resizing to exactly `crop_size` is correct here.

*What is not reproduced.* `EvalVideoTransform` takes `num_views_per_clip` crops
sliding along the long axis and the evaluation loop averages their softmax
outputs. This processor takes a single centre crop. Multi-view aggregation is
the caller's job; `video_io.aggregate_predictions` implements the reference
combination rule.

Compatible with transformers 4.x and 5.x. In particular it does not declare a
`VideosKwargs` subclass as `valid_kwargs`: on transformers 5 an empty TypedDict
subclass loses the field defaults and `preprocess` then fails validation with
`StrictDataclassFieldValidationError`. The inherited `valid_kwargs = VideosKwargs`
is correct on both majors.
"""

import torch

from transformers.image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    PILImageResampling,
    SizeDict,
)
from transformers.processing_utils import Unpack, VideosKwargs
from transformers.video_processing_utils import BaseVideoProcessor


class VJEPA21VideoProcessor(BaseVideoProcessor):
    resample = PILImageResampling.BILINEAR
    image_mean = IMAGENET_DEFAULT_MEAN
    image_std = IMAGENET_DEFAULT_STD
    size = {"shortest_edge": 384}
    crop_size = {"height": 384, "width": 384}
    do_resize = True
    do_rescale = True
    do_center_crop = True
    do_normalize = True
    model_input_names = ["pixel_values_videos"]

    def __init__(self, **kwargs: Unpack[VideosKwargs]):
        crop_size = kwargs.get("crop_size", self.crop_size)
        if isinstance(crop_size, int):
            crop_size = {"height": crop_size, "width": crop_size}
        elif not isinstance(crop_size, dict) or "height" not in crop_size:
            raise ValueError("crop_size must be an integer or a dict with a 'height' key")

        kwargs["crop_size"] = crop_size
        # The short side is resized to the crop size, matching EvalVideoTransform.
        kwargs.setdefault("size", {"shortest_edge": crop_size["height"]})
        super().__init__(**kwargs)

    def resize(
        self,
        image: "torch.Tensor",
        size: SizeDict,
        resample=None,
        antialias: bool = False,
        **kwargs,
    ) -> "torch.Tensor":
        """Resize without antialiasing, matching the reference pipeline.

        The reference transform resizes with `cv2.resize(..., cv2.INTER_LINEAR)`
        (`src/datasets/utils/video/functional.py::resize_clip`), which does not
        antialias. torchvision antialiases by default, and the two agree to
        within one 8-bit level only when upscaling: downscaling a 1080p frame to
        384 they differ by up to 150/255, mean 33/255 — a visibly different image
        fed to a model that never saw antialiased inputs during training.
        """
        return super().resize(image, size, resample=resample, antialias=antialias, **kwargs)


__all__ = ["VJEPA21VideoProcessor"]