"""
Verify VideoPreprocessor (inference) == preprocess_frames_batch (training)
for identical input frames. Training parity: both paths must produce the same tensor.
"""
import cv2
import numpy as np
import torch
from badas.core.preprocessing import (
    VideoPreprocessor,
    preprocess_frames_batch,
    preprocess_single_frame,
)


def make_random_frames(n=16, h=256, w=256, seed=42):
    """Reproducible random RGB uint8 frames."""
    rng = np.random.default_rng(seed)
    return [rng.integers(0, 256, (h, w, 3), dtype=np.uint8) for _ in range(n)]


def test_preprocess_frames_batch_vs_single():
    """preprocess_frames_batch == stacked preprocess_single_frame."""
    frames_rgb = make_random_frames(n=4)
    batch = preprocess_frames_batch(frames_rgb, img_size=256)
    stacked = torch.stack([preprocess_single_frame(f, 256) for f in frames_rgb])
    assert torch.allclose(batch, stacked, atol=1e-5), \
        f"Max diff: {(batch - stacked).abs().max().item()}"


def test_preprocess_window_matches_batch():
    """VideoPreprocessor.preprocess_window(bgr) == preprocess_frames_batch(rgb)."""
    frames_rgb = make_random_frames(n=16, h=480, w=640)  # non-square source
    frames_bgr = [cv2.cvtColor(f, cv2.COLOR_RGB2BGR) for f in frames_rgb]

    training = preprocess_frames_batch(frames_rgb, img_size=256)     # (16, 3, 256, 256)
    p = VideoPreprocessor(img_size=256, num_frames=16, device='cpu')
    inference = p.preprocess_window(frames_bgr).squeeze(0)            # (16, 3, 256, 256)

    assert training.shape == inference.shape
    assert torch.allclose(training, inference, atol=1e-5), \
        f"Max diff: {(training - inference).abs().max().item()}"


def test_add_frame_incremental_matches_batch():
    """add_frame() one at a time == preprocess_window() on same frames."""
    frames_rgb = make_random_frames(n=16, h=480, w=640)
    frames_bgr = [cv2.cvtColor(f, cv2.COLOR_RGB2BGR) for f in frames_rgb]

    p_batch = VideoPreprocessor(img_size=256, num_frames=16, device='cpu')
    batch_t = p_batch.preprocess_window(frames_bgr)

    p_incr = VideoPreprocessor(img_size=256, num_frames=16, device='cpu')
    for f in frames_bgr:
        p_incr.add_frame(f)
    incr_t = p_incr.get_tensor()

    assert torch.allclose(batch_t, incr_t, atol=1e-5), \
        f"Max diff: {(batch_t - incr_t).abs().max().item()}"


def test_rolling_buffer_holds_last_window():
    """After >16 frames, buffer holds the last 16 (matches batch on same 16)."""
    all_rgb = make_random_frames(n=20, h=256, w=256)
    all_bgr = [cv2.cvtColor(f, cv2.COLOR_RGB2BGR) for f in all_rgb]

    # Batch on the last 16
    p_batch = VideoPreprocessor(img_size=256, num_frames=16, device='cpu')
    batch_t = p_batch.preprocess_window(all_bgr[4:])   # frames 4..19

    # Incremental: feed all 20, buffer should hold last 16
    p_incr = VideoPreprocessor(img_size=256, num_frames=16, device='cpu')
    for f in all_bgr:
        p_incr.add_frame(f)
    incr_t = p_incr.get_tensor()

    assert torch.allclose(batch_t, incr_t, atol=1e-5), \
        f"Rolling buffer wrong. Max diff: {(batch_t - incr_t).abs().max().item()}"
