"""Parity tests vs torchvision for both backends, all interp×antialias combos, ragged inputs.

Run locally from the repo root with the package on the path:
    PYTHONPATH=torch-ext pytest tests/ -q
CUDA is required (Triton); tests skip on CPU.
"""

import pytest
import torch
import torchvision.transforms.v2.functional as tvF
from torchvision.transforms import InterpolationMode

from kernel_image_resize import resize_normalize


_TV_INTERP = {"bilinear": InterpolationMode.BILINEAR, "bicubic": InterpolationMode.BICUBIC}

MEAN = [0.48145466, 0.4578275, 0.40821073]
STD = [0.26862954, 0.26130258, 0.27577711]
RESCALE = 1.0 / 255.0


def _ragged_images(n, device, min_res=384, max_res=1024, seed=0):
    g = torch.Generator(device="cpu").manual_seed(seed)
    images = []
    for _ in range(n):
        h = int(torch.randint(min_res, max_res + 1, (1,), generator=g).item())
        w = int(torch.randint(min_res, max_res + 1, (1,), generator=g).item())
        images.append(torch.randint(0, 256, (3, h, w), generator=g, dtype=torch.uint8).to(device))
    return images


def _torchvision_reference(images, out_h, out_w, interp, antialias):
    mode = _TV_INTERP[interp]
    mean = torch.tensor(MEAN, device=images[0].device).view(3, 1, 1)
    std = torch.tensor(STD, device=images[0].device).view(3, 1, 1)
    outs = []
    for img in images:
        r = tvF.resize(img.float(), [out_h, out_w], interpolation=mode, antialias=antialias)
        outs.append((r * RESCALE - mean) / std)
    return torch.stack(outs)


@pytest.mark.kernels_ci
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel needs CUDA")
@pytest.mark.parametrize("backend", ["fused", "separable"])
@pytest.mark.parametrize("interp,antialias", [("bilinear", False), ("bilinear", True), ("bicubic", False), ("bicubic", True)])
def test_parity_vs_torchvision(backend, interp, antialias):
    device = torch.device("cuda")
    images = _ragged_images(8, device)
    out_h = out_w = 384
    got = resize_normalize(
        images, (out_h, out_w), MEAN, STD, RESCALE, resample=interp, antialias=antialias, backend=backend
    )
    ref = _torchvision_reference(images, out_h, out_w, interp, antialias)
    max_abs = (got - ref).abs().max().item()
    assert max_abs < 3e-3, f"{backend}/{interp}/aa={antialias}: max|Δ|={max_abs:.2e}"


@pytest.mark.kernels_ci
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel needs CUDA")
def test_stacked_tensor_input():
    device = torch.device("cuda")
    images = torch.randint(0, 256, (4, 3, 512, 512), dtype=torch.uint8, device=device)
    got = resize_normalize(images, 224, MEAN, STD, RESCALE, resample="bicubic", antialias=True)
    assert got.shape == (4, 3, 224, 224)


def _shortest_edge_crop_reference(images, shortest_edge, crop, interp, antialias):
    mode = _TV_INTERP[interp]
    mean = torch.tensor(MEAN, device=images[0].device).view(3, 1, 1)
    std = torch.tensor(STD, device=images[0].device).view(3, 1, 1)
    outs = []
    for img in images:
        in_h, in_w = img.shape[1], img.shape[2]
        if in_h <= in_w:
            rh, rw = shortest_edge, int(in_w * shortest_edge / in_h)
        else:
            rh, rw = int(in_h * shortest_edge / in_w), shortest_edge
        r = tvF.resize(img.float(), [rh, rw], interpolation=mode, antialias=antialias)
        r = tvF.center_crop(r, [crop, crop])
        outs.append((r * RESCALE - mean) / std)
    return torch.stack(outs)


@pytest.mark.kernels_ci
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel needs CUDA")
@pytest.mark.parametrize("interp,antialias", [("bilinear", True), ("bicubic", True)])
def test_shortest_edge_crop_parity(interp, antialias):
    device = torch.device("cuda")
    images = _ragged_images(8, device)
    shortest_edge, crop = 256, 224
    got = resize_normalize(
        images, shortest_edge, MEAN, STD, RESCALE, resample=interp, antialias=antialias,
        crop_size=(crop, crop), resize_mode="shortest_edge",
    )
    ref = _shortest_edge_crop_reference(images, shortest_edge, crop, interp, antialias)
    assert got.shape == (8, 3, crop, crop)
    max_abs = (got - ref).abs().max().item()
    assert max_abs < 3e-3, f"shortest_edge+crop {interp}/aa={antialias}: max|Δ|={max_abs:.2e}"


@pytest.mark.kernels_ci
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel needs CUDA")
def test_fused_matches_separable():
    device = torch.device("cuda")
    images = _ragged_images(6, device)
    common = dict(size=(256, 256), image_mean=MEAN, image_std=STD, rescale_factor=RESCALE, resample="bicubic", antialias=True)
    fused = resize_normalize(images, backend="fused", **common)
    separable = resize_normalize(images, backend="separable", **common)
    max_abs = (fused - separable).abs().max().item()
    assert max_abs < 3e-3, f"fused vs separable: max|Δ|={max_abs:.2e}"