innr 0.6.2

//! SIMD Correctness Tests: Differential Testing Approach
//!
//! Strategy: Model each SIMD operation in pure scalar Rust, then verify
//! the SIMD implementation matches across random inputs.
//!
//! Based on Cryspen's approach: https://cryspen.com/post/specify-rust-simd/
//! They verified 384 x86 intrinsics + 181 AArch64 intrinsics this way.

#![allow(clippy::float_cmp)]

use innr::{cosine, dot, fast_cosine_dispatch, l2_distance, l2_distance_squared};

// =============================================================================
// Reference Implementations (Pure Scalar)
// =============================================================================

/// Reference dot product - simple, obviously correct
fn ref_dot(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let mut sum = 0.0f32;
    for i in 0..a.len() {
        sum += a[i] * b[i];
    }
    sum
}

/// Reference L2 squared distance
fn ref_l2_squared(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let mut sum = 0.0f32;
    for i in 0..a.len() {
        let d = a[i] - b[i];
        sum += d * d;
    }
    sum
}

/// Reference L2 distance
fn ref_l2(a: &[f32], b: &[f32]) -> f32 {
    ref_l2_squared(a, b).sqrt()
}

/// Reference cosine similarity
fn ref_cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot = ref_dot(a, b);
    let norm_a = ref_l2_squared(a, &vec![0.0; a.len()]).sqrt();
    let norm_b = ref_l2_squared(b, &vec![0.0; b.len()]).sqrt();
    if norm_a < f32::EPSILON || norm_b < f32::EPSILON {
        0.0
    } else {
        dot / (norm_a * norm_b)
    }
}

// =============================================================================
// Test Helpers
// =============================================================================

/// Generate deterministic test vectors
fn test_vec(dim: usize, seed: u64) -> Vec<f32> {
    (0..dim)
        .map(|i| {
            let x = (seed.wrapping_mul(31).wrapping_add(i as u64 * 17)) as f32;
            (x * 0.001).sin()
        })
        .collect()
}

/// Check if two floats are approximately equal
fn approx_eq(a: f32, b: f32, rel_eps: f32) -> bool {
    let diff = (a - b).abs();
    let max_val = a.abs().max(b.abs()).max(1.0);
    diff < rel_eps * max_val
}

// =============================================================================
// Differential Tests
//
// Tolerances are 1e-4 for large dimensions and 1e-5 for small dimensions.
// The gap comes from FMA (fused multiply-add) rounding: SIMD paths use
// hardware FMA which rounds once per multiply-add, while the scalar
// reference rounds twice (once for multiply, once for add). Over hundreds
// of accumulations the rounding difference compounds, requiring the
// looser 1e-4 tolerance for high-dimensional vectors.
// =============================================================================

#[test]
fn simd_correctness_dot_small_dims() {
    // Test dimensions that exercise different SIMD code paths
    // AVX2: 8 floats/register, AVX-512: 16 floats/register
    for dim in [1, 2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] {
        for seed in 0..10 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);

            let simd_result = dot(&a, &b);
            let ref_result = ref_dot(&a, &b);

            assert!(
                approx_eq(simd_result, ref_result, 1e-5),
                "dot mismatch at dim={}, seed={}: simd={}, ref={}",
                dim,
                seed,
                simd_result,
                ref_result
            );
        }
    }
}

#[test]
fn simd_correctness_dot_large_dims() {
    // Typical embedding dimensions
    for dim in [64, 128, 256, 384, 512, 768, 1024, 1536] {
        for seed in 0..5 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);

            let simd_result = dot(&a, &b);
            let ref_result = ref_dot(&a, &b);

            assert!(
                approx_eq(simd_result, ref_result, 1e-4),
                "dot mismatch at dim={}: simd={}, ref={}",
                dim,
                simd_result,
                ref_result
            );
        }
    }
}

#[test]
fn simd_correctness_l2_squared() {
    for dim in [1, 7, 8, 15, 16, 31, 32, 64, 128, 384, 768] {
        for seed in 0..5 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);

            let simd_result = l2_distance_squared(&a, &b);
            let ref_result = ref_l2_squared(&a, &b);

            assert!(
                approx_eq(simd_result, ref_result, 1e-4),
                "l2_squared mismatch at dim={}: simd={}, ref={}",
                dim,
                simd_result,
                ref_result
            );
        }
    }
}

#[test]
fn simd_correctness_l2() {
    for dim in [1, 8, 16, 32, 64, 128, 384, 768] {
        for seed in 0..5 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);

            let simd_result = l2_distance(&a, &b);
            let ref_result = ref_l2(&a, &b);

            assert!(
                approx_eq(simd_result, ref_result, 1e-4),
                "l2 mismatch at dim={}: simd={}, ref={}",
                dim,
                simd_result,
                ref_result
            );
        }
    }
}

#[test]
fn simd_correctness_cosine() {
    for dim in [1, 8, 16, 32, 64, 128, 384, 768] {
        for seed in 0..5 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);

            let simd_result = cosine(&a, &b);
            let ref_result = ref_cosine(&a, &b);

            assert!(
                approx_eq(simd_result, ref_result, 1e-4),
                "cosine mismatch at dim={}: simd={}, ref={}",
                dim,
                simd_result,
                ref_result
            );
        }
    }
}

#[test]
fn simd_correctness_fast_cosine_dispatch() {
    for dim in [1, 8, 16, 32, 64, 128, 384, 768] {
        for seed in 0..5 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);

            let dispatch_result = fast_cosine_dispatch(&a, &b);
            let ref_result = ref_cosine(&a, &b);

            // fast_cosine_dispatch uses rsqrt approximation, so wider tolerance
            assert!(
                approx_eq(dispatch_result, ref_result, 1e-2),
                "fast_cosine_dispatch mismatch at dim={}: dispatch={}, ref={}",
                dim,
                dispatch_result,
                ref_result
            );
        }
    }
}

// =============================================================================
// Edge Cases
// =============================================================================

#[test]
fn simd_correctness_edge_cases() {
    // Empty vectors
    assert_eq!(dot(&[], &[]), 0.0);

    // Single element
    assert_eq!(dot(&[2.0], &[3.0]), 6.0);

    // Zeros
    let zeros = vec![0.0f32; 100];
    let ones = vec![1.0f32; 100];
    assert_eq!(dot(&zeros, &ones), 0.0);
    assert_eq!(l2_distance_squared(&zeros, &zeros), 0.0);

    // Identical vectors
    let v = test_vec(64, 42);
    assert!(approx_eq(l2_distance(&v, &v), 0.0, 1e-6));
    assert!(approx_eq(cosine(&v, &v), 1.0, 1e-5));
}

#[test]
fn simd_correctness_special_values() {
    // Very small values (underflow potential)
    let small: Vec<f32> = (0..64).map(|i| 1e-20 * (i as f32 + 1.0)).collect();
    let ref_result = ref_dot(&small, &small);
    let simd_result = dot(&small, &small);
    // Both should be very small, just check same order of magnitude
    assert!(
        (simd_result - ref_result).abs() < 1e-30,
        "small value mismatch"
    );

    // Large values (overflow potential)
    let large: Vec<f32> = (0..64).map(|i| 1e10 * (i as f32 + 1.0)).collect();
    let ref_result = ref_dot(&large, &large);
    let simd_result = dot(&large, &large);
    assert!(
        approx_eq(simd_result, ref_result, 1e-3),
        "large value mismatch: simd={}, ref={}",
        simd_result,
        ref_result
    );

    // Mixed signs
    let mixed: Vec<f32> = (0..64)
        .map(|i| if i % 2 == 0 { 1.0 } else { -1.0 })
        .collect();
    let ref_result = ref_dot(&mixed, &mixed);
    let simd_result = dot(&mixed, &mixed);
    assert_eq!(simd_result, ref_result);
}

// =============================================================================
// Invariant Tests
// =============================================================================

#[test]
fn simd_invariant_dot_commutative() {
    for dim in [32, 64, 128, 384] {
        let a = test_vec(dim, 1);
        let b = test_vec(dim, 2);
        assert!(
            approx_eq(dot(&a, &b), dot(&b, &a), 1e-6),
            "dot should be commutative"
        );
    }
}

#[test]
fn simd_invariant_l2_symmetric() {
    for dim in [32, 64, 128, 384] {
        let a = test_vec(dim, 1);
        let b = test_vec(dim, 2);
        assert!(
            approx_eq(l2_distance(&a, &b), l2_distance(&b, &a), 1e-6),
            "l2 should be symmetric"
        );
    }
}

#[test]
fn simd_invariant_l2_nonnegative() {
    for dim in [32, 64, 128, 384] {
        let a = test_vec(dim, 1);
        let b = test_vec(dim, 2);
        assert!(l2_distance(&a, &b) >= 0.0, "l2 should be non-negative");
        assert!(
            l2_distance_squared(&a, &b) >= 0.0,
            "l2_squared should be non-negative"
        );
    }
}

#[test]
fn simd_invariant_cosine_range() {
    for dim in [32, 64, 128, 384] {
        for seed in 0..10 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 100);
            let sim = cosine(&a, &b);
            assert!(
                (-1.0 - 1e-5..=1.0 + 1e-5).contains(&sim),
                "cosine should be in [-1, 1], got {}",
                sim
            );
        }
    }
}

// =============================================================================
// L1 distance: differential vs scalar reference across dispatch boundaries.
// Previously the only oracle test was NEON-gated in-file; a tail bug in the
// AVX2/AVX-512 kernels would have passed the x86 CI job.
// =============================================================================

fn ref_l1(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| (x - y).abs()).sum()
}

#[test]
fn simd_correctness_l1_across_boundaries() {
    for dim in [
        1, 2, 3, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 256, 768,
    ] {
        for seed in 0..5 {
            let a = test_vec(dim, seed);
            let b = test_vec(dim, seed + 1000);
            let simd_result = innr::l1_distance(&a, &b);
            let ref_result = ref_l1(&a, &b);
            assert!(
                approx_eq(simd_result, ref_result, 1e-4),
                "l1 mismatch at dim={dim}, seed={seed}: simd={simd_result}, ref={ref_result}"
            );
        }
    }
}

// =============================================================================
// mixed f32 x u8 dot: exact recovery with identity quantization params.
// The SIMD kernels were previously tested only through quantization-scale
// tolerances, which would mask a sub-tolerance kernel bug.
// =============================================================================

#[test]
fn simd_correctness_mixed_dot_exact() {
    use innr::scalar::{asymmetric_dot_u8, quantize_u8, QuantizationParams};
    // alpha = 255, offset = 0 makes dequantization the identity:
    // result = (alpha/255) * mixed_dot + offset * query_sum = mixed_dot.
    let params = QuantizationParams {
        alpha: 255.0,
        offset: 0.0,
    };
    for dim in [8, 16, 31, 32, 33, 64, 65, 128] {
        for seed in 0..5_u64 {
            // u8 corpus values and small-integer f32 queries: every product
            // and partial sum is exactly representable in f32, so the SIMD
            // kernel must match the reference bit-for-bit.
            let corpus: Vec<f32> = (0..dim)
                .map(|i| ((i as u64 * 31 + seed * 7) % 256) as f32)
                .collect();
            let query: Vec<f32> = (0..dim)
                .map(|i| ((i as u64 * 13 + seed * 3) % 8) as f32)
                .collect();
            let qz = quantize_u8(&corpus, &params);
            let expect: f32 = query.iter().zip(corpus.iter()).map(|(x, y)| x * y).sum();
            let got = asymmetric_dot_u8(&query, &qz, &params);
            assert_eq!(got, expect, "mixed dot mismatch at dim={dim}, seed={seed}");
        }
    }
}