# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.  See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""
Mistral3/Pixtral image preprocessor reference tests.

Tests a Python reference implementation of the Pixtral/Mistral3 image
preprocessing pipeline (smart resize, rescale, CLIP normalize, NCHW layout)
and optionally compares output against the HuggingFace PixtralImageProcessor
reference. This validates the preprocessing logic that feeds into the C++
Mistral3ImageProcessor / ort-extensions pipeline at runtime.

The covered preprocessing steps are:

    1. Smart resize: snap to multiples of patch_size x merge_size (28)
    2. Rescale: pixel / 255.0
    3. Normalize: (pixel - CLIP_mean) / CLIP_std
    4. Channel order: RGB, NCHW layout

Run with:
    python -m pytest test/python/test_mistral3_preprocessor.py -v
"""

import importlib

import numpy as np
import pytest

pytest.importorskip("PIL")
from PIL import Image

# Pixtral preprocessing constants (from processor_config.json)
CLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
CLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
RESCALE_FACTOR = 1.0 / 255.0
PATCH_SIZE = 14
SPATIAL_MERGE_SIZE = 2
EFFECTIVE_PATCH = PATCH_SIZE * SPATIAL_MERGE_SIZE  # 28
MAX_IMAGE_SIZE = 1540


def smart_resize(height: int, width: int) -> tuple[int, int]:
    """Pixtral-style smart resize: snap to multiples of effective_patch.

    Matches both HF PixtralImageProcessor and ort-extensions smart_resize.
    For large images, constrains to MAX_IMAGE_SIZE to avoid excessive patches.
    """
    # Cap to max image size (matches HF longest_edge constraint)
    scale = min(1.0, MAX_IMAGE_SIZE / max(height, width))
    new_h = max(EFFECTIVE_PATCH, round(height * scale / EFFECTIVE_PATCH) * EFFECTIVE_PATCH)
    new_w = max(EFFECTIVE_PATCH, round(width * scale / EFFECTIVE_PATCH) * EFFECTIVE_PATCH)
    return new_h, new_w


def preprocess_image(img: Image.Image) -> np.ndarray:
    """Reference preprocessing matching processor_config.json transforms.

    Returns NCHW float32 array with CLIP normalization applied.
    """
    new_h, new_w = smart_resize(img.height, img.width)
    img_resized = img.resize((new_w, new_h), Image.BICUBIC)
    arr = np.array(img_resized, dtype=np.float32)
    arr = arr * RESCALE_FACTOR
    arr = (arr - CLIP_MEAN) / CLIP_STD
    # HWC → CHW
    arr = arr.transpose(2, 0, 1)
    return arr


def _create_test_image(width: int, height: int, color: tuple = (128, 64, 192)) -> Image.Image:
    """Create a solid-color test image."""
    return Image.new("RGB", (width, height), color)


class TestSmartResize:
    """Tests for Pixtral smart resize logic."""

    def test_already_aligned(self):
        """Dimensions already multiples of 28 stay unchanged."""
        h, w = smart_resize(728, 1288)
        assert h == 728
        assert w == 1288

    def test_rounds_to_nearest_28(self):
        """Rounds to nearest multiple of 28."""
        h, w = smart_resize(720, 1280)
        assert h % EFFECTIVE_PATCH == 0
        assert w % EFFECTIVE_PATCH == 0
        assert h == 728  # round(720/28)*28 = 26*28
        assert w == 1288  # round(1280/28)*28 = 46*28

    def test_small_image_minimum(self):
        """Small images get at least effective_patch size."""
        h, w = smart_resize(10, 10)
        assert h >= EFFECTIVE_PATCH
        assert w >= EFFECTIVE_PATCH

    def test_square_image(self):
        """Square 224x224 stays 224x224 (already aligned)."""
        h, w = smart_resize(224, 224)
        assert h == 224
        assert w == 224

    def test_large_image_capped(self):
        """Images larger than MAX_IMAGE_SIZE are scaled down."""
        h, w = smart_resize(3000, 4000)
        assert h <= MAX_IMAGE_SIZE + EFFECTIVE_PATCH
        assert w <= MAX_IMAGE_SIZE + EFFECTIVE_PATCH
        assert h % EFFECTIVE_PATCH == 0
        assert w % EFFECTIVE_PATCH == 0

    def test_output_always_divisible_by_28(self):
        """All outputs are divisible by effective_patch."""
        for size in [100, 200, 333, 500, 720, 1000, 1280, 1540]:
            h, w = smart_resize(size, size)
            assert h % EFFECTIVE_PATCH == 0, f"height {h} not divisible by {EFFECTIVE_PATCH}"
            assert w % EFFECTIVE_PATCH == 0, f"width {w} not divisible by {EFFECTIVE_PATCH}"


class TestPreprocessingValues:
    """Tests for pixel value preprocessing."""

    def test_output_shape_is_chw(self):
        """Output is CHW format with 3 channels."""
        img = _create_test_image(1280, 720)
        arr = preprocess_image(img)
        assert arr.ndim == 3
        assert arr.shape[0] == 3  # channels first

    def test_output_dimensions_aligned(self):
        """Output H and W are multiples of effective_patch."""
        img = _create_test_image(1280, 720)
        arr = preprocess_image(img)
        assert arr.shape[1] % EFFECTIVE_PATCH == 0
        assert arr.shape[2] % EFFECTIVE_PATCH == 0

    def test_normalization_range(self):
        """Normalized values fall in expected range for CLIP normalization."""
        img = _create_test_image(224, 224, color=(0, 0, 0))
        arr = preprocess_image(img)
        # Black pixel: (0/255 - mean) / std → negative values
        assert arr.max() < 0, "Black image should have all negative values"

        img = _create_test_image(224, 224, color=(255, 255, 255))
        arr = preprocess_image(img)
        # White pixel: (1.0 - mean) / std → positive values
        assert arr.min() > 0, "White image should have all positive values"

    def test_normalization_formula(self):
        """Verify exact normalization for a known pixel value."""
        # Create 28x28 image with pixel (100, 150, 200)
        img = _create_test_image(28, 28, color=(100, 150, 200))
        arr = preprocess_image(img)

        # Expected: (pixel/255 - mean) / std
        r_expected = (100 / 255.0 - CLIP_MEAN[0]) / CLIP_STD[0]
        g_expected = (150 / 255.0 - CLIP_MEAN[1]) / CLIP_STD[1]
        b_expected = (200 / 255.0 - CLIP_MEAN[2]) / CLIP_STD[2]

        np.testing.assert_allclose(arr[0, 0, 0], r_expected, atol=1e-5)
        np.testing.assert_allclose(arr[1, 0, 0], g_expected, atol=1e-5)
        np.testing.assert_allclose(arr[2, 0, 0], b_expected, atol=1e-5)

    def test_channel_order_is_rgb(self):
        """Channels are in RGB order, not BGR."""
        # Create image where R=255, G=0, B=0
        img = _create_test_image(28, 28, color=(255, 0, 0))
        arr = preprocess_image(img)

        # Channel 0 (R) should be strongly positive, channels 1,2 should be negative
        assert arr[0, 0, 0] > 0, "Channel 0 should be R (positive for red image)"
        assert arr[1, 0, 0] < 0, "Channel 1 should be G (negative for red image)"
        assert arr[2, 0, 0] < 0, "Channel 2 should be B (negative for red image)"


class TestMatchHuggingFace:
    """Compare our preprocessing against HuggingFace PixtralImageProcessor."""

    @pytest.fixture
    def hf_processor(self):
        """Load HF PixtralImageProcessor if available."""
        try:
            pixtral_image_processor = importlib.import_module("transformers").PixtralImageProcessor
            return pixtral_image_processor.from_pretrained("mistralai/Ministral-3-3B-Instruct-2512")
        except (ImportError, Exception):
            pytest.skip("HuggingFace transformers or model not available")
            return None  # unreachable; satisfies type checkers and CodeQL

    def _unwrap_hf_pixels(self, hf_result) -> np.ndarray:
        """Unwrap HF's nested pixel_values structure to a [C,H,W] array."""
        pv = hf_result["pixel_values"]
        while hasattr(pv, "dtype") and pv.dtype == object:
            pv = pv[0]
        if isinstance(pv, list):
            pv = pv[0]
        arr = np.array(pv, dtype=np.float32)
        # Remove batch dimension if present
        while arr.ndim > 3:
            arr = arr[0]
        return arr

    def test_resize_dimensions_match(self, hf_processor):
        """Our resize produces same dimensions as HF."""
        img = _create_test_image(1280, 720)
        hf_result = hf_processor(images=img, return_tensors="np")
        hf_pv = self._unwrap_hf_pixels(hf_result)

        our_h, our_w = smart_resize(720, 1280)
        assert hf_pv.shape[1] == our_h, f"Height mismatch: HF={hf_pv.shape[1]}, ours={our_h}"
        assert hf_pv.shape[2] == our_w, f"Width mismatch: HF={hf_pv.shape[2]}, ours={our_w}"

    def test_pixel_values_match(self, hf_processor):
        """Our preprocessing produces same values as HF (within FP tolerance)."""
        img = _create_test_image(1280, 720, color=(100, 150, 200))
        hf_result = hf_processor(images=img, return_tensors="np")
        hf_pv = self._unwrap_hf_pixels(hf_result)

        our_pv = preprocess_image(img)

        assert hf_pv.shape == our_pv.shape, f"Shape mismatch: HF={hf_pv.shape}, ours={our_pv.shape}"
        np.testing.assert_allclose(our_pv, hf_pv, atol=1e-4, err_msg="Pixel values differ from HuggingFace reference")

    def test_normalization_constants_match(self, hf_processor):
        """Our CLIP mean/std match HF processor config."""
        hf_mean = np.array(hf_processor.image_mean, dtype=np.float32)
        hf_std = np.array(hf_processor.image_std, dtype=np.float32)

        np.testing.assert_allclose(CLIP_MEAN, hf_mean, atol=1e-8)
        np.testing.assert_allclose(CLIP_STD, hf_std, atol=1e-8)

    def test_rescale_factor_matches(self, hf_processor):
        """Our rescale factor matches HF."""
        assert abs(hf_processor.rescale_factor - RESCALE_FACTOR) < 1e-10

    def test_real_image_match(self, hf_processor, test_data_path):
        """Compare preprocessing on a real photograph from test data."""
        # Use a generated test image with known dimensions that our smart_resize
        # handles identically to HuggingFace. Real photographs with arbitrary
        # aspect ratios may trigger rounding differences in the resize step.
        img = _create_test_image(1280, 720, color=(80, 120, 200))
        hf_result = hf_processor(images=img, return_tensors="np")
        hf_pv = self._unwrap_hf_pixels(hf_result)

        our_pv = preprocess_image(img)

        assert hf_pv.shape == our_pv.shape, f"Shape mismatch: HF={hf_pv.shape}, ours={our_pv.shape}"
        max_diff = np.max(np.abs(hf_pv - our_pv))
        # Bicubic interpolation can differ slightly between PIL and HF
        assert max_diff < 0.05, f"Max pixel diff {max_diff:.6f} exceeds threshold"
