"""End-to-end parity: video processor + model against the reference pipeline.

Sections 2 and 3 of the model card verify the model given an input tensor, and
the preprocessing is verified separately against the documented constants. This
file closes the loop: one clip goes through the reference
`Resize -> crop -> ClipToTensor -> Normalize` and through
`VJEPA21VideoProcessor`, and both tensors are then pushed through the model.

    pip install opencv-python-headless
    VJEPA2_REPO=/path/to/vjepa2 python -m pytest test_end_to_end_pipeline.py -s -q

Add VJEPA21_CKPT to measure the feature-level consequence on real weights:

    VJEPA2_REPO=/path/to/vjepa2 VJEPA21_CKPT=apiantonio/vjepa2.1-vit-base-384 \
        python -m pytest test_end_to_end_pipeline.py -s -q
"""

from __future__ import annotations

import importlib
import os
import sys
import types

import numpy as np
import pytest
import torch

PORT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
VJEPA2_REPO = os.environ.get("VJEPA2_REPO", "")
CKPT = os.environ.get("VJEPA21_CKPT", "")


def _processor(**kwargs):
    name = "_vjepa21_port_under_test"
    if name not in sys.modules:
        pkg = types.ModuleType(name)
        pkg.__path__ = [PORT_DIR]
        sys.modules[name] = pkg
    mod = importlib.import_module(f"{name}.video_processing_vjepa21")
    return mod.VJEPA21VideoProcessor(**kwargs)


def _reference_transform(clip, crop=384):
    """Reference `EvalVideoTransform`, centre crop instead of the sliding views."""
    repo = os.path.abspath(VJEPA2_REPO)
    if repo not in sys.path:
        sys.path.insert(0, repo)
    import src.datasets.utils.video.transforms as video_transforms
    import src.datasets.utils.video.volume_transforms as volume_transforms

    to_tensor = video_transforms.Compose(
        [
            volume_transforms.ClipToTensor(),
            video_transforms.Normalize(
                mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)
            ),
        ]
    )
    buf = np.array(video_transforms.Resize(crop, interpolation="bilinear")(clip))
    _, h, w, _ = buf.shape
    start = (max(h, w) - crop) // 2
    view = buf[:, :, start : start + crop, :] if w >= h else buf[:, start : start + crop, :, :]
    return to_tensor(list(view)).float()  # (C, T, H, W)


def _clip(height, width, frames=4, seed=0):
    rng = np.random.default_rng(seed)
    return [rng.integers(0, 255, (height, width, 3), dtype=np.uint8) for _ in range(frames)]


# --- 1. antialiasing: no reference repository needed -------------------------


def test_resize_defaults_to_no_antialiasing():
    """Regression guard for the resize backend.

    This used to be run on `torch.zeros(1, 3, 1080, 1920)`, which cannot fail:
    resizing a constant image gives the same result with and without
    antialiasing, so `torch.equal(default, forced)` held regardless of the
    default. It needs real content, and it needs to show that the antialiased
    path is genuinely different — otherwise the guard is still vacuous.
    """
    processor = _processor()
    torch.manual_seed(0)
    big = torch.rand(1, 3, 1080, 1920)
    size = type(processor.size)(shortest_edge=384)

    default = processor.resize(big, size)
    forced_off = processor.resize(big, size, antialias=False)
    forced_on = processor.resize(big, size, antialias=True)

    assert torch.equal(default, forced_off), "resize() must default to antialias=False"
    gap = (forced_on - default).abs().max().item()
    print(f"\n[antialias] max|Δ| antialiased vs not = {gap:.4f} "
          f"({gap * 255:.1f}/255 on an 8-bit scale)")
    assert gap > 1e-3, (
        "antialiased and non-antialiased downscaling are indistinguishable; "
        "the guard is not testing anything"
    )


def test_antialiasing_gap_is_negligible_when_upscaling():
    """Explains why this only ever mattered for downscaling: on a 240p source
    resized up to 384 the two backends agree, which is how the bug survived."""
    processor = _processor()
    torch.manual_seed(0)
    small = torch.rand(1, 3, 240, 320)
    size = type(processor.size)(shortest_edge=384)
    gap = (
        processor.resize(small, size, antialias=True)
        - processor.resize(small, size, antialias=False)
    ).abs().max().item()
    print(f"[antialias, upscale] max|Δ| = {gap:.6f}")
    assert gap < 1e-5


def test_processor_geometry_is_independent_of_aspect_ratio():
    """Short side to exactly `crop_size`, then a square centre crop."""
    processor = _processor()
    for h, w in ((240, 320), (320, 240), (1080, 1920), (400, 400)):
        pv = processor([_clip(h, w, frames=2)], return_tensors="pt")["pixel_values_videos"]
        assert pv.shape == (1, 2, 3, 384, 384), (h, w, pv.shape)


# --- 2. pixel level against the reference transform --------------------------

RESOLUTIONS = [
    pytest.param(240, 320, id="240p-upscale"),
    pytest.param(480, 858, id="480p-downscale"),
    pytest.param(720, 1280, id="720p-downscale"),
    pytest.param(1080, 1920, id="1080p-downscale"),
]

needs_reference = pytest.mark.skipif(
    not VJEPA2_REPO or not os.path.isdir(os.path.join(VJEPA2_REPO, "src", "datasets")),
    reason="set VJEPA2_REPO to a clone of facebookresearch/vjepa2",
)


@needs_reference
@pytest.mark.parametrize("height,width", RESOLUTIONS)
def test_processor_matches_reference_transform(height, width):
    """The processor must agree with the reference transform to within the 8-bit
    quantisation floor, at every input resolution.

    The floor is one grey level divided by the smallest normalisation std:
    1/255 / 0.229 = 1.75e-02. Anything materially above that means the resize
    backends disagree — the failure mode here is torchvision's default
    antialiasing, which the reference `cv2.INTER_LINEAR` path does not apply and
    which only shows up when downscaling.
    """
    clip = _clip(height, width)
    reference = _reference_transform(clip)
    got = _processor()([clip], return_tensors="pt")["pixel_values_videos"][0]
    got = got.permute(1, 0, 2, 3)  # (T, C, H, W) -> (C, T, H, W)

    assert reference.shape == got.shape
    diff = (reference - got).abs()
    quantisation_floor = 1.0 / 255.0 / 0.229
    print(
        f"\n[{height}p] max|Δ| = {diff.max():.4e}  mean|Δ| = {diff.mean():.4e}  "
        f"rel = {diff.mean() / reference.abs().mean():.4e}  (floor {quantisation_floor:.4e})"
    )
    assert diff.max().item() <= quantisation_floor * 1.05


@needs_reference
def test_antialiased_processor_would_fail_the_reference_comparison():
    """Shows the tolerance above is tight enough to catch the regression it was
    written for, rather than passing either way."""
    clip = _clip(1080, 1920)
    reference = _reference_transform(clip)

    processor = _processor()
    original_resize = type(processor).resize

    def antialiased(self, image, size, **kwargs):
        kwargs["antialias"] = True
        return super(type(self), self).resize(image, size, **kwargs)

    try:
        type(processor).resize = antialiased
        got = processor([clip], return_tensors="pt")["pixel_values_videos"][0]
        got = got.permute(1, 0, 2, 3)
    finally:
        type(processor).resize = original_resize

    gap = (reference - got).abs().max().item()
    print(f"\n[antialiased 1080p] max|Δ| vs reference = {gap:.4e}")
    assert gap > 1.0 / 255.0 / 0.229 * 1.05


# --- 3. feature level --------------------------------------------------------


@needs_reference
@pytest.mark.skipif(not CKPT, reason="set VJEPA21_CKPT to run the model-level comparison")
@torch.no_grad()
@pytest.mark.parametrize("height,width", RESOLUTIONS)
def test_features_match_reference_pipeline(height, width):
    """What the preprocessing residual costs in feature space."""
    from transformers import AutoModel

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained(CKPT, trust_remote_code=True).eval().to(device)

    clip = _clip(height, width, frames=16)
    reference = _reference_transform(clip).unsqueeze(0).to(device)  # (1, C, T, H, W)
    processed = _processor()([clip], return_tensors="pt")["pixel_values_videos"].to(device)

    a = model(pixel_values_videos=reference, skip_predictor=True).last_hidden_state
    b = model(pixel_values_videos=processed, skip_predictor=True).last_hidden_state

    rel = ((a - b).abs().mean() / a.abs().mean()).item()
    cos = torch.nn.functional.cosine_similarity(a.flatten(0, 1), b.flatten(0, 1)).min().item()
    print(f"\n[{height}p features] rel = {rel:.4e}  min cos-sim = {cos:.6f}")
    assert cos > 0.99