//! Golden tests for vision processors.
//!
//! These tests compare Rust preprocessor output against golden outputs
//! generated by HuggingFace transformers to ensure pixel-perfect compatibility.
//!
//! Modes tested:
//! - `llava/` - Standard CLIP processing (llava-hf/* models, no expand-to-square)
//! - `llava_pad/` - Expand-to-square mode (liuhaotian/llava-* models, image_aspect_ratio=pad)
//! - `qwen2_vl/` - Dynamic resolution with smart resize (Qwen/Qwen2-VL-* models)
//! - `qwen3_vl/` - Dynamic resolution with patch_size=16 and [0.5,0.5,0.5] norm (Qwen/Qwen3-VL-* models)
//!
//! To regenerate golden outputs:
//! ```bash
//! python scripts/generate_vision_golden.py
//! ```

use std::{fs::File, io::Read, path::Path};

use ndarray::{Array4, Array5};
use sgl_model_gateway::multimodal::vision::{
    image_processor::ModelSpecificValue, ImagePreProcessor, Llama4VisionProcessor, LlavaProcessor,
    Phi3VisionProcessor, Phi4VisionProcessor, PixtralProcessor, PreProcessorConfig,
    Qwen2VLProcessor, Qwen3VLProcessor,
};

/// Load a numpy .npz file and extract pixel_values
fn load_golden_npz(path: &Path) -> Array4<f32> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    // Read pixel_values array (npz stores arrays without .npy extension in the lookup)
    let reader = npz
        .by_name("pixel_values")
        .expect("Failed to read npz")
        .expect("No pixel_values");

    // Get shape from npy header
    let shape = reader.shape().to_vec();
    assert_eq!(shape.len(), 4, "Expected 4D tensor [B, C, H, W]");

    // Read data as f32 vec
    let data: Vec<f32> = reader.into_vec().expect("Failed to read array");

    // Convert to Array4
    Array4::from_shape_vec(
        (
            shape[0] as usize,
            shape[1] as usize,
            shape[2] as usize,
            shape[3] as usize,
        ),
        data,
    )
    .expect("Shape conversion failed")
}

/// Load preprocessor config from JSON
fn load_config(path: &Path) -> PreProcessorConfig {
    let mut file = File::open(path).expect("Failed to open config");
    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .expect("Failed to read config");
    PreProcessorConfig::from_json(&contents).expect("Failed to parse config")
}

/// Compare two 4D tensors and return max absolute difference
fn max_diff(a: &Array4<f32>, b: &ndarray::ArrayD<f32>) -> f32 {
    assert_eq!(a.shape(), b.shape(), "Shape mismatch");
    // Convert ArrayD to Array4 for comparison
    let b_4d = b
        .clone()
        .into_dimensionality::<ndarray::Ix4>()
        .expect("Expected 4D tensor");
    (a - &b_4d)
        .mapv(|v| v.abs())
        .fold(0.0f32, |acc, &v| acc.max(v))
}

/// Load image_grid_thw from npz file
fn load_golden_grid_thw(path: &Path) -> Vec<u32> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("image_grid_thw")
        .expect("Failed to read npz")
        .expect("No image_grid_thw");

    // Shape not needed, data is flat
    let _shape = reader.shape();

    // Read data as i64 vec (numpy default for int)
    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");

    // Convert to u32
    data.into_iter().map(|v| v as u32).collect()
}

/// Load num_tokens from npz file
fn load_golden_num_tokens(path: &Path) -> usize {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("num_tokens")
        .expect("Failed to read npz")
        .expect("No num_tokens");

    // Read single value as i64
    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");
    data[0] as usize
}

/// Run a golden test for a specific mode and image.
///
/// # Arguments
/// * `mode` - Either "llava" (standard CLIP) or "llava_pad" (expand-to-square mode)
/// * `image_name` - Name of the test image (e.g., "square", "tall", "wide", "small")
fn run_golden_test(mode: &str, image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden").join(mode);
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for {}/{} not found, skipping test",
            mode, image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py");
        return;
    }

    let golden = load_golden_npz(&golden_dir.join(format!("golden_{}.npz", image_name)));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    let image = image::open(&image_path).expect("Failed to open image");

    let processor: Box<dyn ImagePreProcessor> = match mode {
        "llava" => Box::new(LlavaProcessor::new()),
        "llava_pad" => Box::new(LlavaProcessor::new_with_pad()),
        _ => panic!("Unknown test mode: {}", mode),
    };

    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    let diff = max_diff(&golden, &result.pixel_values);
    println!(
        "{} - {} image - Max difference: {:.6}",
        mode, image_name, diff
    );
    println!("Golden shape: {:?}", golden.shape());
    println!("Rust shape: {:?}", result.pixel_values.shape());

    // Allow tolerance for floating point and interpolation algorithm differences
    // Different interpolation implementations (Rust vs Python/PIL) can produce
    // small numerical differences, especially for edge cases like tiny or extreme
    // aspect ratio images
    assert!(diff < 0.1, "Max difference {} exceeds tolerance 0.1", diff);
}

// ============================================================================
// Standard CLIP mode tests (llava-hf/* models, no expand-to-square)
// ============================================================================

#[test]
fn test_llava_golden_square() {
    run_golden_test("llava", "square");
}

#[test]
fn test_llava_golden_tall() {
    run_golden_test("llava", "tall");
}

#[test]
fn test_llava_golden_wide() {
    run_golden_test("llava", "wide");
}

#[test]
fn test_llava_golden_small() {
    run_golden_test("llava", "small");
}

#[test]
fn test_llava_golden_tiny() {
    run_golden_test("llava", "tiny");
}

#[test]
fn test_llava_golden_very_tall() {
    run_golden_test("llava", "very_tall");
}

#[test]
fn test_llava_golden_very_wide() {
    run_golden_test("llava", "very_wide");
}

#[test]
fn test_llava_golden_large() {
    run_golden_test("llava", "large");
}

#[test]
fn test_llava_golden_odd_dims() {
    run_golden_test("llava", "odd_dims");
}

#[test]
fn test_llava_golden_grayscale() {
    run_golden_test("llava", "grayscale");
}

// ============================================================================
// Pad mode tests (liuhaotian/llava-* models, image_aspect_ratio=pad)
// ============================================================================

#[test]
fn test_llava_pad_golden_square() {
    run_golden_test("llava_pad", "square");
}

#[test]
fn test_llava_pad_golden_tall() {
    run_golden_test("llava_pad", "tall");
}

#[test]
fn test_llava_pad_golden_wide() {
    run_golden_test("llava_pad", "wide");
}

#[test]
fn test_llava_pad_golden_small() {
    run_golden_test("llava_pad", "small");
}

#[test]
fn test_llava_pad_golden_tiny() {
    run_golden_test("llava_pad", "tiny");
}

#[test]
fn test_llava_pad_golden_very_tall() {
    run_golden_test("llava_pad", "very_tall");
}

#[test]
fn test_llava_pad_golden_very_wide() {
    run_golden_test("llava_pad", "very_wide");
}

#[test]
fn test_llava_pad_golden_large() {
    run_golden_test("llava_pad", "large");
}

#[test]
fn test_llava_pad_golden_odd_dims() {
    run_golden_test("llava_pad", "odd_dims");
}

#[test]
fn test_llava_pad_golden_grayscale() {
    run_golden_test("llava_pad", "grayscale");
}

// ============================================================================
// Token count tests
// ============================================================================

#[test]
fn test_llava_token_count() {
    let golden_dir = Path::new("tests/fixtures/golden/llava");

    if !golden_dir.exists() {
        eprintln!("Golden test fixtures not found, skipping test");
        return;
    }

    let config = load_config(&golden_dir.join("preprocessor_config.json"));
    let processor = LlavaProcessor::new();

    // LLaVA 1.5 with 336x336 and patch_size=14: (336/14)^2 = 576 tokens
    let tokens = processor.calculate_num_tokens(336, 336, &config);
    assert_eq!(
        tokens, 576,
        "Expected 576 tokens for 336x336 with patch_size=14"
    );
}

// ============================================================================
// Qwen2-VL tests
// ============================================================================

/// Load flattened pixel values from Qwen2-VL npz file.
/// Returns (data, shape) where shape is (num_patches, patch_features).
fn load_golden_qwen2_vl_pixels(path: &Path) -> (Vec<f32>, (usize, usize)) {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("pixel_values")
        .expect("Failed to read npz")
        .expect("No pixel_values");

    let shape = reader.shape().to_vec();
    assert_eq!(shape.len(), 2, "Expected 2D tensor for Qwen2-VL patches");

    let data: Vec<f32> = reader.into_vec().expect("Failed to read array");
    (data, (shape[0] as usize, shape[1] as usize))
}

/// Run a Qwen2-VL golden test for a specific image.
///
/// This test validates:
/// 1. image_grid_thw matches the HuggingFace output
/// 2. num_tokens calculation is correct
/// 3. Pixel values match after reshaping to patch format
fn run_qwen2_vl_golden_test(image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden/qwen2_vl");
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for qwen2_vl/{} not found, skipping test",
            image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py --model qwen2_vl");
        return;
    }

    let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    // Load golden values
    let golden_grid_thw = load_golden_grid_thw(&npz_path);
    let golden_num_tokens = load_golden_num_tokens(&npz_path);
    let (golden_pixels, golden_shape) = load_golden_qwen2_vl_pixels(&npz_path);

    // Process image with our Rust processor
    let image = image::open(&image_path).expect("Failed to open image");
    let processor = Qwen2VLProcessor::from_preprocessor_config(&config);
    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    // Extract image_grid_thw from result
    let rust_grid_thw = match result.model_specific.get("image_grid_thw") {
        Some(ModelSpecificValue::UintTensor { data, shape }) => {
            assert_eq!(shape, &[1, 3], "Expected shape [1, 3] for single image");
            data.clone()
        }
        _ => panic!("Expected image_grid_thw in model_specific"),
    };

    // Compare grid dimensions
    println!(
        "qwen2_vl - {} image - Grid T H W: golden={:?}, rust={:?}",
        image_name, golden_grid_thw, rust_grid_thw
    );
    assert_eq!(
        golden_grid_thw, rust_grid_thw,
        "image_grid_thw mismatch for {}",
        image_name
    );

    // Compare token counts
    let rust_num_tokens = result.num_img_tokens[0];
    println!(
        "qwen2_vl - {} image - Tokens: golden={}, rust={}",
        image_name, golden_num_tokens, rust_num_tokens
    );
    assert_eq!(
        golden_num_tokens, rust_num_tokens,
        "num_tokens mismatch for {}",
        image_name
    );

    // Compare pixel values by reshaping our output to patch format
    let grid_t = rust_grid_thw[0] as usize;
    let grid_h = rust_grid_thw[1] as usize;
    let grid_w = rust_grid_thw[2] as usize;

    // Get the tensor for the first image (batch index 0)
    let pixel_values = &result.pixel_values;
    let tensor_3d_dyn = pixel_values.index_axis(ndarray::Axis(0), 0).to_owned();
    let tensor_3d = tensor_3d_dyn
        .into_dimensionality::<ndarray::Ix3>()
        .expect("Expected 3D tensor for Qwen2-VL");

    // Reshape to patches format
    let rust_patches = processor.reshape_to_patches(&tensor_3d, grid_t, grid_h, grid_w);

    // Verify shapes match
    let expected_num_patches = grid_t * grid_h * grid_w;
    let patch_size = config.get_patch_size(14);
    let temporal_patch_size = config.temporal_patch_size.unwrap_or(2);
    let expected_patch_features = 3 * temporal_patch_size * patch_size * patch_size;

    println!(
        "qwen2_vl - {} image - Patch shape: golden={:?}, rust=({}, {})",
        image_name, golden_shape, expected_num_patches, expected_patch_features
    );
    assert_eq!(
        golden_shape,
        (expected_num_patches, expected_patch_features),
        "Patch shape mismatch"
    );
    assert_eq!(
        rust_patches.len(),
        expected_num_patches * expected_patch_features,
        "Rust patches size mismatch"
    );

    // Compare pixel values
    let max_diff = rust_patches
        .iter()
        .zip(golden_pixels.iter())
        .map(|(r, g)| (r - g).abs())
        .fold(0.0f32, f32::max);

    println!(
        "qwen2_vl - {} image - Max pixel diff: {:.6}",
        image_name, max_diff
    );

    // Allow tolerance for floating point and interpolation differences
    // Different interpolation implementations (Rust vs Python/PIL) can produce
    // small numerical differences, especially for edge cases
    assert!(
        max_diff < 0.1,
        "Max pixel difference {} exceeds tolerance 0.1 for {}",
        max_diff,
        image_name
    );
}

#[test]
fn test_qwen2_vl_golden_square() {
    run_qwen2_vl_golden_test("square");
}

#[test]
fn test_qwen2_vl_golden_tall() {
    run_qwen2_vl_golden_test("tall");
}

#[test]
fn test_qwen2_vl_golden_wide() {
    run_qwen2_vl_golden_test("wide");
}

#[test]
fn test_qwen2_vl_golden_small() {
    run_qwen2_vl_golden_test("small");
}

#[test]
fn test_qwen2_vl_golden_tiny() {
    run_qwen2_vl_golden_test("tiny");
}

#[test]
fn test_qwen2_vl_golden_very_tall() {
    run_qwen2_vl_golden_test("very_tall");
}

#[test]
fn test_qwen2_vl_golden_very_wide() {
    run_qwen2_vl_golden_test("very_wide");
}

#[test]
fn test_qwen2_vl_golden_large() {
    run_qwen2_vl_golden_test("large");
}

#[test]
fn test_qwen2_vl_golden_odd_dims() {
    run_qwen2_vl_golden_test("odd_dims");
}

#[test]
fn test_qwen2_vl_golden_grayscale() {
    run_qwen2_vl_golden_test("grayscale");
}

// ============================================================================
// Qwen3-VL tests
// ============================================================================

/// Run a Qwen3-VL golden test for a specific image.
///
/// This test validates:
/// 1. image_grid_thw matches the HuggingFace output
/// 2. num_tokens calculation is correct
/// 3. Pixel values match after reshaping to patch format
///
/// Key differences from Qwen2-VL:
/// - patch_size: 16 (vs 14)
/// - factor: 32 (vs 28)
/// - normalization: [0.5, 0.5, 0.5] (vs CLIP)
fn run_qwen3_vl_golden_test(image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden/qwen3_vl");
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for qwen3_vl/{} not found, skipping test",
            image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py --model qwen3_vl");
        return;
    }

    let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    // Load golden values
    let golden_grid_thw = load_golden_grid_thw(&npz_path);
    let golden_num_tokens = load_golden_num_tokens(&npz_path);
    let (golden_pixels, golden_shape) = load_golden_qwen2_vl_pixels(&npz_path);

    // Process image with our Rust processor
    let image = image::open(&image_path).expect("Failed to open image");
    let processor = Qwen3VLProcessor::from_preprocessor_config(&config);
    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    // Extract image_grid_thw from result
    let rust_grid_thw = match result.model_specific.get("image_grid_thw") {
        Some(ModelSpecificValue::UintTensor { data, shape }) => {
            assert_eq!(shape, &[1, 3], "Expected shape [1, 3] for single image");
            data.clone()
        }
        _ => panic!("Expected image_grid_thw in model_specific"),
    };

    // Compare grid dimensions
    println!(
        "qwen3_vl - {} image - Grid T H W: golden={:?}, rust={:?}",
        image_name, golden_grid_thw, rust_grid_thw
    );
    assert_eq!(
        golden_grid_thw, rust_grid_thw,
        "image_grid_thw mismatch for {}",
        image_name
    );

    // Compare token counts
    let rust_num_tokens = result.num_img_tokens[0];
    println!(
        "qwen3_vl - {} image - Tokens: golden={}, rust={}",
        image_name, golden_num_tokens, rust_num_tokens
    );
    assert_eq!(
        golden_num_tokens, rust_num_tokens,
        "num_tokens mismatch for {}",
        image_name
    );

    // Compare pixel values by reshaping our output to patch format
    let grid_t = rust_grid_thw[0] as usize;
    let grid_h = rust_grid_thw[1] as usize;
    let grid_w = rust_grid_thw[2] as usize;

    // Get the tensor for the first image (batch index 0)
    let pixel_values = &result.pixel_values;
    let tensor_3d_dyn = pixel_values.index_axis(ndarray::Axis(0), 0).to_owned();
    let tensor_3d = tensor_3d_dyn
        .into_dimensionality::<ndarray::Ix3>()
        .expect("Expected 3D tensor for Qwen3-VL");

    // Reshape to patches format
    let rust_patches = processor.reshape_to_patches(&tensor_3d, grid_t, grid_h, grid_w);

    // Verify shapes match (Qwen3-VL has patch_size=16)
    let expected_num_patches = grid_t * grid_h * grid_w;
    let patch_size = config.get_patch_size(16);
    let temporal_patch_size = config.temporal_patch_size.unwrap_or(2);
    let expected_patch_features = 3 * temporal_patch_size * patch_size * patch_size;

    println!(
        "qwen3_vl - {} image - Patch shape: golden={:?}, rust=({}, {})",
        image_name, golden_shape, expected_num_patches, expected_patch_features
    );
    assert_eq!(
        golden_shape,
        (expected_num_patches, expected_patch_features),
        "Patch shape mismatch"
    );
    assert_eq!(
        rust_patches.len(),
        expected_num_patches * expected_patch_features,
        "Rust patches size mismatch"
    );

    // Compare pixel values
    let max_diff = rust_patches
        .iter()
        .zip(golden_pixels.iter())
        .map(|(r, g)| (r - g).abs())
        .fold(0.0f32, f32::max);

    println!(
        "qwen3_vl - {} image - Max pixel diff: {:.6}",
        image_name, max_diff
    );

    // Allow tolerance for floating point and interpolation differences
    // Max diff is ~0.03 due to resize interpolation differences between Rust and HuggingFace
    assert!(
        max_diff < 0.05,
        "Max pixel difference {} exceeds tolerance 0.05 for {}",
        max_diff,
        image_name
    );
}

#[test]
fn test_qwen3_vl_golden_square() {
    run_qwen3_vl_golden_test("square");
}

#[test]
fn test_qwen3_vl_golden_tall() {
    run_qwen3_vl_golden_test("tall");
}

#[test]
fn test_qwen3_vl_golden_wide() {
    run_qwen3_vl_golden_test("wide");
}

#[test]
fn test_qwen3_vl_golden_small() {
    run_qwen3_vl_golden_test("small");
}

#[test]
fn test_qwen3_vl_golden_tiny() {
    run_qwen3_vl_golden_test("tiny");
}

#[test]
fn test_qwen3_vl_golden_very_tall() {
    run_qwen3_vl_golden_test("very_tall");
}

#[test]
fn test_qwen3_vl_golden_very_wide() {
    run_qwen3_vl_golden_test("very_wide");
}

#[test]
fn test_qwen3_vl_golden_large() {
    run_qwen3_vl_golden_test("large");
}

#[test]
fn test_qwen3_vl_golden_odd_dims() {
    run_qwen3_vl_golden_test("odd_dims");
}

#[test]
fn test_qwen3_vl_golden_grayscale() {
    run_qwen3_vl_golden_test("grayscale");
}

// ============================================================================
// Phi3-Vision tests
// ============================================================================

/// Load a 5D numpy .npz file for Phi3-Vision (batch, num_crops+1, C, H, W)
fn load_golden_npz_5d(path: &Path) -> Array5<f32> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("pixel_values")
        .expect("Failed to read npz")
        .expect("No pixel_values");

    let shape = reader.shape().to_vec();
    assert_eq!(shape.len(), 5, "Expected 5D tensor [B, N, C, H, W]");

    let data: Vec<f32> = reader.into_vec().expect("Failed to read array");

    Array5::from_shape_vec(
        (
            shape[0] as usize,
            shape[1] as usize,
            shape[2] as usize,
            shape[3] as usize,
            shape[4] as usize,
        ),
        data,
    )
    .expect("Shape conversion failed")
}

/// Load image_sizes from Phi3-Vision npz file (2D tensor [batch, 2])
fn load_phi3_image_sizes(path: &Path) -> Vec<(u32, u32)> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("image_sizes")
        .expect("Failed to read npz")
        .expect("No image_sizes");

    let shape = reader.shape().to_vec();
    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");

    // Reshape to pairs
    let num_images = shape[0] as usize;
    (0..num_images)
        .map(|i| (data[i * 2] as u32, data[i * 2 + 1] as u32))
        .collect()
}

/// Load num_img_tokens from Phi3-Vision npz file
fn load_phi3_num_img_tokens(path: &Path) -> Vec<usize> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("num_img_tokens")
        .expect("Failed to read npz")
        .expect("No num_img_tokens");

    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");
    data.into_iter().map(|v| v as usize).collect()
}

/// Compare two 5D tensors and return max absolute difference
fn max_diff_5d(a: &Array5<f32>, b: &Array5<f32>) -> f32 {
    assert_eq!(a.shape(), b.shape(), "Shape mismatch");
    (a - b).mapv(|v| v.abs()).fold(0.0f32, |acc, &v| acc.max(v))
}

/// Find the location and value of max difference between two 5D tensors
#[allow(dead_code)]
fn find_max_diff_location_5d(
    golden: &Array5<f32>,
    rust: &Array5<f32>,
    image_name: &str,
) -> (f32, (usize, usize, usize, usize, usize)) {
    assert_eq!(golden.shape(), rust.shape(), "Shape mismatch");
    let diff = (golden - rust).mapv(|v| v.abs());
    let mut max_diff = 0.0f32;
    let mut max_pos = (0, 0, 0, 0, 0);

    // Find per-tile max differences
    for b in 0..golden.shape()[0] {
        for t in 0..golden.shape()[1] {
            let tile_diff = diff.slice(ndarray::s![b, t, .., .., ..]);
            let tile_max = tile_diff.fold(0.0f32, |acc, &v| acc.max(v));

            if tile_max > 0.1 {
                let golden_tile = golden.slice(ndarray::s![b, t, .., .., ..]);
                let rust_tile = rust.slice(ndarray::s![b, t, .., .., ..]);
                println!(
                    "  {} tile {}: diff={:.4}, golden_range=[{:.4}, {:.4}], rust_range=[{:.4}, {:.4}]",
                    image_name, t, tile_max,
                    golden_tile.fold(f32::MAX, |a, &v| a.min(v)),
                    golden_tile.fold(f32::MIN, |a, &v| a.max(v)),
                    rust_tile.fold(f32::MAX, |a, &v| a.min(v)),
                    rust_tile.fold(f32::MIN, |a, &v| a.max(v))
                );
            }

            if tile_max > max_diff {
                max_diff = tile_max;
                // Find exact position
                for c in 0..golden.shape()[2] {
                    for h in 0..golden.shape()[3] {
                        for w in 0..golden.shape()[4] {
                            if diff[[b, t, c, h, w]] == max_diff {
                                max_pos = (b, t, c, h, w);
                            }
                        }
                    }
                }
            }
        }
    }

    (max_diff, max_pos)
}

/// Run a Phi3-Vision golden test for a specific image.
///
/// This test validates:
/// 1. Output shape is [1, num_crops+1, 3, 336, 336]
/// 2. image_sizes matches HuggingFace output
/// 3. num_img_tokens matches HuggingFace output
/// 4. Pixel values match within tolerance
fn run_phi3_vision_golden_test(image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden/phi3_vision");
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for phi3_vision/{} not found, skipping test",
            image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py --model phi3_vision");
        return;
    }

    let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    // Load golden values
    let golden_pixels = load_golden_npz_5d(&npz_path);
    let golden_image_sizes = load_phi3_image_sizes(&npz_path);
    let golden_num_tokens = load_phi3_num_img_tokens(&npz_path);

    // Process image with our Rust processor
    let image = image::open(&image_path).expect("Failed to open image");
    let processor = Phi3VisionProcessor::from_preprocessor_config(&config);
    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    // Check output shape
    let rust_shape = result.pixel_values.shape();
    let golden_shape = golden_pixels.shape();
    println!(
        "phi3_vision - {} image - Shape: golden={:?}, rust={:?}",
        image_name, golden_shape, rust_shape
    );
    assert_eq!(
        rust_shape, golden_shape,
        "Shape mismatch for phi3_vision/{}",
        image_name
    );

    // Check image_sizes
    // Note: HuggingFace returns [h, w], we store as (w, h) but model_specific stores (h, w)
    let rust_image_sizes: Vec<(u32, u32)> = match result.model_specific.get("image_sizes") {
        Some(ModelSpecificValue::UintTensor { data, shape }) => {
            let num_images = shape[0];
            (0..num_images)
                .map(|i| (data[i * 2], data[i * 2 + 1]))
                .collect()
        }
        _ => panic!("Expected image_sizes in model_specific"),
    };

    println!(
        "phi3_vision - {} image - Image sizes (h, w): golden={:?}, rust={:?}",
        image_name, golden_image_sizes, rust_image_sizes
    );
    assert_eq!(
        golden_image_sizes, rust_image_sizes,
        "image_sizes mismatch for {}",
        image_name
    );

    // Check num_img_tokens
    println!(
        "phi3_vision - {} image - Num tokens: golden={:?}, rust={:?}",
        image_name, golden_num_tokens, result.num_img_tokens
    );
    assert_eq!(
        golden_num_tokens, result.num_img_tokens,
        "num_img_tokens mismatch for {}",
        image_name
    );

    // Compare pixel values
    // Convert rust ArrayD to Array5 for comparison
    let rust_pixels = result
        .pixel_values
        .clone()
        .into_dimensionality::<ndarray::Ix5>()
        .expect("Failed to convert to Ix5");

    let pixel_diff = max_diff_5d(&golden_pixels, &rust_pixels);
    println!(
        "phi3_vision - {} image - Max pixel diff: {:.6}",
        image_name, pixel_diff
    );

    // If there's a large difference, print detailed info
    if pixel_diff > 0.1 {
        let (max_diff, max_pos) =
            find_max_diff_location_5d(&golden_pixels, &rust_pixels, image_name);
        println!(
            "phi3_vision - {} image - Max diff {:.4} at position {:?}",
            image_name, max_diff, max_pos
        );
        let (b, t, c, h, w) = max_pos;
        println!(
            "  golden value: {:.4}, rust value: {:.4}",
            golden_pixels[[b, t, c, h, w]],
            rust_pixels[[b, t, c, h, w]]
        );
    }

    // Allow tolerance for floating point and interpolation differences
    // Using bicubic for global image and bilinear for HD resize to match HuggingFace.
    assert!(
        pixel_diff < 0.08,
        "Max pixel difference {} exceeds tolerance 0.08 for {}",
        pixel_diff,
        image_name
    );
}

#[test]
fn test_phi3_vision_golden_square() {
    run_phi3_vision_golden_test("square");
}

#[test]
fn test_phi3_vision_golden_tall() {
    run_phi3_vision_golden_test("tall");
}

#[test]
fn test_phi3_vision_golden_wide() {
    run_phi3_vision_golden_test("wide");
}

#[test]
fn test_phi3_vision_golden_small() {
    run_phi3_vision_golden_test("small");
}

#[test]
fn test_phi3_vision_golden_tiny() {
    run_phi3_vision_golden_test("tiny");
}

#[test]
fn test_phi3_vision_golden_very_tall() {
    run_phi3_vision_golden_test("very_tall");
}

#[test]
fn test_phi3_vision_golden_very_wide() {
    run_phi3_vision_golden_test("very_wide");
}

#[test]
fn test_phi3_vision_golden_large() {
    run_phi3_vision_golden_test("large");
}

#[test]
fn test_phi3_vision_golden_odd_dims() {
    run_phi3_vision_golden_test("odd_dims");
}

#[test]
fn test_phi3_vision_golden_grayscale() {
    run_phi3_vision_golden_test("grayscale");
}

// ============================================================================
// Phi4-Vision tests
// ============================================================================

/// Load num_img_tokens from Phi4-Vision npz file
fn load_phi4_num_img_tokens(path: &Path) -> Vec<usize> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("num_img_tokens")
        .expect("Failed to read npz")
        .expect("No num_img_tokens");

    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");
    data.into_iter().map(|v| v as usize).collect()
}

/// Load image_sizes from Phi4-Vision npz file (2D tensor [batch, 2])
fn load_phi4_image_sizes(path: &Path) -> Vec<(u32, u32)> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("image_sizes")
        .expect("Failed to read npz")
        .expect("No image_sizes");

    let shape = reader.shape().to_vec();
    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");

    // Reshape to pairs
    let num_images = shape[0] as usize;
    (0..num_images)
        .map(|i| (data[i * 2] as u32, data[i * 2 + 1] as u32))
        .collect()
}

/// Run a Phi4-Vision golden test for a specific image.
///
/// This test validates:
/// 1. Output shape is [1, num_crops+1, 3, 448, 448] (note: 448 base resolution)
/// 2. image_sizes matches HuggingFace output
/// 3. num_img_tokens matches HuggingFace output
/// 4. Pixel values match within tolerance
///
/// Key differences from Phi3-Vision:
/// - Base resolution: 448 (vs 336)
/// - Normalization: [0.5, 0.5, 0.5] (vs CLIP)
/// - Default dynamic_hd: 36 (vs 16)
fn run_phi4_vision_golden_test(image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden/phi4_vision");
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for phi4_vision/{} not found, skipping test",
            image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py --model phi4_vision");
        return;
    }

    let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    // Load golden values
    let golden_pixels = load_golden_npz_5d(&npz_path);
    let golden_image_sizes = load_phi4_image_sizes(&npz_path);
    let golden_num_tokens = load_phi4_num_img_tokens(&npz_path);

    // Process image with our Rust processor
    let image = image::open(&image_path).expect("Failed to open image");
    let processor = Phi4VisionProcessor::from_preprocessor_config(&config);
    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    // Check output shape
    let rust_shape = result.pixel_values.shape();
    let golden_shape = golden_pixels.shape();
    println!(
        "phi4_vision - {} image - Shape: golden={:?}, rust={:?}",
        image_name, golden_shape, rust_shape
    );
    assert_eq!(
        rust_shape, golden_shape,
        "Shape mismatch for phi4_vision/{}",
        image_name
    );

    // Check image_sizes
    let rust_image_sizes: Vec<(u32, u32)> = match result.model_specific.get("image_sizes") {
        Some(ModelSpecificValue::UintTensor { data, shape }) => {
            let num_images = shape[0];
            (0..num_images)
                .map(|i| (data[i * 2], data[i * 2 + 1]))
                .collect()
        }
        _ => panic!("Expected image_sizes in model_specific"),
    };

    println!(
        "phi4_vision - {} image - Image sizes (h, w): golden={:?}, rust={:?}",
        image_name, golden_image_sizes, rust_image_sizes
    );
    assert_eq!(
        golden_image_sizes, rust_image_sizes,
        "image_sizes mismatch for {}",
        image_name
    );

    // Check num_img_tokens
    println!(
        "phi4_vision - {} image - Num tokens: golden={:?}, rust={:?}",
        image_name, golden_num_tokens, result.num_img_tokens
    );
    assert_eq!(
        golden_num_tokens, result.num_img_tokens,
        "num_img_tokens mismatch for {}",
        image_name
    );

    // Compare pixel values
    let rust_pixels = result
        .pixel_values
        .clone()
        .into_dimensionality::<ndarray::Ix5>()
        .expect("Failed to convert to Ix5");

    let pixel_diff = max_diff_5d(&golden_pixels, &rust_pixels);
    println!(
        "phi4_vision - {} image - Max pixel diff: {:.6}",
        image_name, pixel_diff
    );

    // If there's a large difference, print detailed info
    if pixel_diff > 0.1 {
        let (max_diff, max_pos) =
            find_max_diff_location_5d(&golden_pixels, &rust_pixels, image_name);
        println!(
            "phi4_vision - {} image - Max diff {:.4} at position {:?}",
            image_name, max_diff, max_pos
        );
        let (b, t, c, h, w) = max_pos;
        println!(
            "  golden value: {:.4}, rust value: {:.4}",
            golden_pixels[[b, t, c, h, w]],
            rust_pixels[[b, t, c, h, w]]
        );
    }

    // Allow tolerance for floating point and interpolation differences
    // Using bilinear for HD resize and bicubic for global image to match HuggingFace.
    assert!(
        pixel_diff < 0.05,
        "Max pixel difference {} exceeds tolerance 0.05 for {}",
        pixel_diff,
        image_name
    );
}

#[test]
fn test_phi4_vision_golden_square() {
    run_phi4_vision_golden_test("square");
}

#[test]
fn test_phi4_vision_golden_tall() {
    run_phi4_vision_golden_test("tall");
}

#[test]
fn test_phi4_vision_golden_wide() {
    run_phi4_vision_golden_test("wide");
}

#[test]
fn test_phi4_vision_golden_small() {
    run_phi4_vision_golden_test("small");
}

#[test]
fn test_phi4_vision_golden_tiny() {
    run_phi4_vision_golden_test("tiny");
}

#[test]
fn test_phi4_vision_golden_very_tall() {
    run_phi4_vision_golden_test("very_tall");
}

#[test]
fn test_phi4_vision_golden_very_wide() {
    run_phi4_vision_golden_test("very_wide");
}

#[test]
fn test_phi4_vision_golden_large() {
    run_phi4_vision_golden_test("large");
}

#[test]
fn test_phi4_vision_golden_odd_dims() {
    run_phi4_vision_golden_test("odd_dims");
}

#[test]
fn test_phi4_vision_golden_grayscale() {
    run_phi4_vision_golden_test("grayscale");
}

// ============================================================================
// LLaMA 4 Vision tests
// ============================================================================

/// Load aspect_ratios from npz file for LLaMA 4
fn load_llama4_aspect_ratios(path: &Path) -> Vec<(u32, u32)> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("aspect_ratios")
        .expect("Failed to read npz")
        .expect("No aspect_ratios");

    let shape = reader.shape().to_vec();

    // Read data as i64 vec (numpy default for int)
    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");

    // Convert to Vec<(u32, u32)>
    let num_images = shape[0] as usize;
    (0..num_images)
        .map(|i| (data[i * 2] as u32, data[i * 2 + 1] as u32))
        .collect()
}

/// Load pixel_values for LLaMA 4 Vision (3D: [num_tiles, C, H, W])
fn load_llama4_pixels(path: &Path) -> (Vec<f32>, Vec<usize>) {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("pixel_values")
        .expect("Failed to read npz")
        .expect("No pixel_values");

    let shape: Vec<usize> = reader.shape().iter().map(|&s| s as usize).collect();
    let data: Vec<f32> = reader.into_vec().expect("Failed to read array");

    (data, shape)
}

/// Run a LLaMA 4 Vision golden test for a specific image.
///
/// This test validates:
/// 1. Output shape matches (batch, num_tiles, 3, 336, 336)
/// 2. aspect_ratios match (h_tiles, w_tiles)
/// 3. Pixel values match HuggingFace output
/// 4. Token count is correct
///
/// LLaMA 4 Vision processing:
/// - Tile size: 336x336
/// - Max patches: 16 (default)
/// - Normalization: [0.5, 0.5, 0.5] mean/std
/// - Global tile added when num_tiles > 1
fn run_llama4_vision_golden_test(image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden/llama4_vision");
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for llama4_vision/{} not found, skipping test",
            image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py --model llama4_vision");
        return;
    }

    let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    // Load golden values
    let (golden_pixels, golden_shape) = load_llama4_pixels(&npz_path);
    let golden_aspect_ratios = load_llama4_aspect_ratios(&npz_path);
    let golden_num_tokens = load_golden_num_tokens(&npz_path);

    // Process image with our Rust processor
    let image = image::open(&image_path).expect("Failed to open image");
    let processor = Llama4VisionProcessor::from_preprocessor_config(&config);
    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    // Check aspect_ratios
    let rust_aspect_ratios: Vec<(u32, u32)> = match result.model_specific.get("aspect_ratios") {
        Some(ModelSpecificValue::UintTensor { data, shape }) => {
            let num_images = shape[0];
            (0..num_images)
                .map(|i| (data[i * 2], data[i * 2 + 1]))
                .collect()
        }
        _ => panic!("Expected aspect_ratios in model_specific"),
    };

    println!(
        "llama4_vision - {} image - Aspect ratios: golden={:?}, rust={:?}",
        image_name, golden_aspect_ratios, rust_aspect_ratios
    );
    assert_eq!(
        golden_aspect_ratios, rust_aspect_ratios,
        "aspect_ratios mismatch for {}",
        image_name
    );

    // Check num_tokens
    let rust_num_tokens = result.num_img_tokens[0];
    println!(
        "llama4_vision - {} image - Tokens: golden={}, rust={}",
        image_name, golden_num_tokens, rust_num_tokens
    );
    assert_eq!(
        golden_num_tokens, rust_num_tokens,
        "num_tokens mismatch for {}",
        image_name
    );

    // Check output shape - HuggingFace outputs (num_tiles, 3, 336, 336) without batch
    // Our Rust outputs (batch, num_tiles, 3, 336, 336) with batch dimension
    let rust_shape = result.pixel_values.shape();
    println!(
        "llama4_vision - {} image - Shape: golden={:?}, rust={:?}",
        image_name, golden_shape, rust_shape
    );

    // HuggingFace returns without batch dim, we add batch=1
    assert!(
        rust_shape[0] == 1,
        "Expected batch dim to be 1, got {}",
        rust_shape[0]
    );
    assert!(
        rust_shape[1] >= golden_shape[0],
        "Expected at least {} tiles, got {}",
        golden_shape[0],
        rust_shape[1]
    );

    // Compare pixel values
    let rust_pixels = result.pixel_values_flat();
    let num_golden_elements: usize = golden_shape.iter().product();

    // Find the max difference for the actual tiles (not padding)
    let mut max_diff = 0.0f32;
    for i in 0..num_golden_elements {
        let diff = (rust_pixels[i] - golden_pixels[i]).abs();
        max_diff = max_diff.max(diff);
    }

    println!(
        "llama4_vision - {} image - Max pixel diff: {:.6}",
        image_name, max_diff
    );

    // Allow tolerance for floating point and interpolation differences
    // LLaMA 4 uses bfloat16 internally which may cause small differences
    assert!(
        max_diff < 0.03,
        "Max pixel difference {} exceeds tolerance 0.03 for {}",
        max_diff,
        image_name
    );
}

#[test]
fn test_llama4_vision_golden_square() {
    run_llama4_vision_golden_test("square");
}

#[test]
fn test_llama4_vision_golden_tall() {
    run_llama4_vision_golden_test("tall");
}

#[test]
fn test_llama4_vision_golden_wide() {
    run_llama4_vision_golden_test("wide");
}

#[test]
fn test_llama4_vision_golden_small() {
    run_llama4_vision_golden_test("small");
}

#[test]
fn test_llama4_vision_golden_tiny() {
    run_llama4_vision_golden_test("tiny");
}

#[test]
fn test_llama4_vision_golden_very_tall() {
    run_llama4_vision_golden_test("very_tall");
}

#[test]
fn test_llama4_vision_golden_very_wide() {
    run_llama4_vision_golden_test("very_wide");
}

#[test]
fn test_llama4_vision_golden_large() {
    run_llama4_vision_golden_test("large");
}

#[test]
fn test_llama4_vision_golden_odd_dims() {
    run_llama4_vision_golden_test("odd_dims");
}

#[test]
fn test_llama4_vision_golden_grayscale() {
    run_llama4_vision_golden_test("grayscale");
}

// ============================================================================
// Pixtral/Mistral3 Vision tests
// ============================================================================

/// Load image_sizes from npz file for Pixtral
fn load_pixtral_image_sizes(path: &Path) -> Vec<(usize, usize)> {
    let file = File::open(path).expect("Failed to open golden file");
    let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz");

    let reader = npz
        .by_name("image_sizes")
        .expect("Failed to read npz")
        .expect("No image_sizes");

    let shape = reader.shape().to_vec();

    // Read data as i64 vec (numpy default for int)
    let data: Vec<i64> = reader.into_vec().expect("Failed to read array");

    // Convert to Vec<(usize, usize)>
    let num_images = shape[0] as usize;
    (0..num_images)
        .map(|i| (data[i * 2] as usize, data[i * 2 + 1] as usize))
        .collect()
}

/// Run a Pixtral golden test for a specific image.
///
/// This test validates:
/// 1. Output shape matches (batch, 3, H, W)
/// 2. image_sizes match
/// 3. Pixel values match HuggingFace output
/// 4. Token count is correct
///
/// Pixtral processing:
/// - Longest edge: 1024 (default)
/// - Patch size: 16
/// - Normalization: CLIP mean/std
/// - No tiling - single output per image
fn run_pixtral_golden_test(image_name: &str) {
    let golden_dir = Path::new("tests/fixtures/golden/pixtral");
    let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));

    if !golden_dir.exists() || !image_path.exists() {
        eprintln!(
            "Golden test fixtures for pixtral/{} not found, skipping test",
            image_name
        );
        eprintln!("Run: python scripts/generate_vision_golden.py --model pixtral");
        return;
    }

    let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
    let config = load_config(&golden_dir.join("preprocessor_config.json"));

    // Load golden values
    let golden_pixels = load_golden_npz(&npz_path);
    let golden_shape: Vec<usize> = golden_pixels.shape().to_vec();
    let golden_image_sizes = load_pixtral_image_sizes(&npz_path);
    let golden_num_tokens = load_golden_num_tokens(&npz_path);

    // Process image with our Rust processor
    let image = image::open(&image_path).expect("Failed to open image");
    let processor = PixtralProcessor::from_preprocessor_config(&config);
    let result = processor
        .preprocess(&[image], &config)
        .expect("Processing failed");

    // Check image_sizes from model_specific
    let rust_image_sizes: Vec<(usize, usize)> = match result.model_specific.get("image_sizes") {
        Some(ModelSpecificValue::IntTensor { data, shape }) => {
            let num_images = shape[0];
            (0..num_images)
                .map(|i| (data[i * 2] as usize, data[i * 2 + 1] as usize))
                .collect()
        }
        _ => panic!("Expected image_sizes in model_specific"),
    };

    println!(
        "pixtral - {} image - Image sizes: golden={:?}, rust={:?}",
        image_name, golden_image_sizes, rust_image_sizes
    );
    assert_eq!(
        golden_image_sizes, rust_image_sizes,
        "image_sizes mismatch for {}",
        image_name
    );

    // Check num_tokens
    let rust_num_tokens = result.num_img_tokens[0];
    println!(
        "pixtral - {} image - Tokens: golden={}, rust={}",
        image_name, golden_num_tokens, rust_num_tokens
    );
    assert_eq!(
        golden_num_tokens, rust_num_tokens,
        "num_tokens mismatch for {}",
        image_name
    );

    // Check output shape
    let rust_shape = result.pixel_values.shape();
    println!(
        "pixtral - {} image - Shape: golden={:?}, rust={:?}",
        image_name, golden_shape, rust_shape
    );

    // Pixtral outputs [batch, C, H, W] with padding to max size in batch
    // Single image should match golden shape exactly
    assert_eq!(rust_shape[0], 1, "Expected batch dim to be 1");
    assert_eq!(rust_shape[1], golden_shape[1], "Channel mismatch");
    assert!(
        rust_shape[2] >= golden_shape[2],
        "Height {} < golden height {}",
        rust_shape[2],
        golden_shape[2]
    );
    assert!(
        rust_shape[3] >= golden_shape[3],
        "Width {} < golden width {}",
        rust_shape[3],
        golden_shape[3]
    );

    // Compare pixel values - only compare the actual image region, not padding
    let rust_pixels = result.pixel_values_flat();
    let golden_pixels_flat: Vec<f32> = golden_pixels.iter().copied().collect();

    // Calculate indices for the actual image region (not padding)
    let h = golden_shape[2];
    let w = golden_shape[3];
    let rust_w = rust_shape[3];

    let mut max_diff = 0.0f32;
    for c in 0..3 {
        for y in 0..h {
            for x in 0..w {
                let golden_idx = c * h * w + y * w + x;
                let rust_idx = c * rust_shape[2] * rust_w + y * rust_w + x;
                let diff = (rust_pixels[rust_idx] - golden_pixels_flat[golden_idx]).abs();
                max_diff = max_diff.max(diff);
            }
        }
    }

    println!(
        "pixtral - {} image - Max pixel diff: {:.6}",
        image_name, max_diff
    );

    // Allow tolerance for bicubic interpolation differences between PIL and Rust image library
    // Pixtral uses bicubic which has larger differences than bilinear
    assert!(
        max_diff < 0.06,
        "Max pixel difference {} exceeds tolerance 0.06 for {}",
        max_diff,
        image_name
    );
}

#[test]
fn test_pixtral_golden_square() {
    run_pixtral_golden_test("square");
}

#[test]
fn test_pixtral_golden_tall() {
    run_pixtral_golden_test("tall");
}

#[test]
fn test_pixtral_golden_wide() {
    run_pixtral_golden_test("wide");
}

#[test]
fn test_pixtral_golden_small() {
    run_pixtral_golden_test("small");
}

#[test]
fn test_pixtral_golden_tiny() {
    run_pixtral_golden_test("tiny");
}

#[test]
fn test_pixtral_golden_very_tall() {
    run_pixtral_golden_test("very_tall");
}

#[test]
fn test_pixtral_golden_very_wide() {
    run_pixtral_golden_test("very_wide");
}

#[test]
fn test_pixtral_golden_large() {
    run_pixtral_golden_test("large");
}

#[test]
fn test_pixtral_golden_odd_dims() {
    run_pixtral_golden_test("odd_dims");
}

#[test]
fn test_pixtral_golden_grayscale() {
    run_pixtral_golden_test("grayscale");
}