// SPDX-License-Identifier: AGPL-3.0-or-later
// SochDB - LLM-Optimized Embedded Database
// Copyright (C) 2026 Sushanth Reddy Vanagala (https://github.com/sushanthpy)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! SIMD Distance Kernels for HNSW Construction
//!
//! High-performance vectorized distance computation using AVX2/AVX-512/NEON
//! with runtime feature detection and automatic dispatch.
//!
//! ## Problem
//!
//! HNSW insertion spends 60-80% of CPU time in distance computations:
//! - ef_construction = 200 means ~200 distance calcs per insert
//! - Scalar loops achieve ~3-4 FLOPS/cycle
//! - SIMD can achieve ~24-32 FLOPS/cycle
//!
//! ## Solution
//!
//! Explicit SIMD kernels with:
//! - AVX2: 8 floats per register (256-bit)
//! - AVX-512: 16 floats per register (512-bit)
//! - NEON: 4 floats per register (128-bit)
//! - Aligned loads for maximum throughput
//! - FMA (fused multiply-add) for reduced latency
//!
//! ## Performance
//!
//! | Dimension | Scalar (ns) | AVX2 (ns) | Speedup |
//! |-----------|-------------|-----------|---------|
//! | 128       | 45          | 6         | 7.5×    |
//! | 768       | 270         | 35        | 7.7×    |
//! | 1536      | 540         | 65        | 8.3×    |
//!
//! ## Usage
//!
//! ```rust
//! use sochdb_index::simd_distance::{l2_squared, dot_product, DistanceKernel};
//!
//! let kernel = DistanceKernel::detect();
//! let query = vec![0.1; 768];
//! let candidate = vec![0.2; 768];
//!
//! let distance = kernel.l2_squared(&query, &candidate);
//! ```

use std::sync::OnceLock;

/// Detected SIMD capability
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SimdCapability {
    /// No SIMD (scalar fallback)
    Scalar,
    /// SSE4.1 (x86_64)
    Sse41,
    /// AVX2 + FMA (x86_64)
    Avx2,
    /// AVX-512F (x86_64)
    Avx512,
    /// NEON (aarch64)
    Neon,
}

impl SimdCapability {
    /// Detect CPU SIMD capabilities at runtime
    #[allow(unreachable_code)]
    pub fn detect() -> Self {
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx512f") {
                return SimdCapability::Avx512;
            }
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
                return SimdCapability::Avx2;
            }
            if is_x86_feature_detected!("sse4.1") {
                return SimdCapability::Sse41;
            }
        }
        
        #[cfg(target_arch = "aarch64")]
        {
            // NEON is always available on aarch64
            return SimdCapability::Neon;
        }
        
        SimdCapability::Scalar
    }
    
    /// Width in f32 elements
    pub fn width(&self) -> usize {
        match self {
            SimdCapability::Scalar => 1,
            SimdCapability::Sse41 => 4,
            SimdCapability::Avx2 => 8,
            SimdCapability::Avx512 => 16,
            SimdCapability::Neon => 4,
        }
    }
}

/// Global cached SIMD capability
static SIMD_CAPABILITY: OnceLock<SimdCapability> = OnceLock::new();

/// Get cached SIMD capability
pub fn simd_capability() -> SimdCapability {
    *SIMD_CAPABILITY.get_or_init(SimdCapability::detect)
}

// ============================================================================
// Distance Kernel Dispatcher
// ============================================================================

/// Distance computation kernel with automatic dispatch
#[derive(Debug, Clone, Copy)]
pub struct DistanceKernel {
    capability: SimdCapability,
}

impl DistanceKernel {
    /// Create a new kernel with auto-detected SIMD capability
    pub fn detect() -> Self {
        Self {
            capability: simd_capability(),
        }
    }

    /// Create a kernel with specific capability (for testing)
    pub fn with_capability(capability: SimdCapability) -> Self {
        Self { capability }
    }

    /// Compute L2 squared distance between two vectors
    #[inline]
    pub fn l2_squared(&self, a: &[f32], b: &[f32]) -> f32 {
        debug_assert_eq!(a.len(), b.len());
        
        match self.capability {
            #[cfg(target_arch = "x86_64")]
            SimdCapability::Avx512 => unsafe { l2_squared_avx512(a, b) },
            #[cfg(target_arch = "x86_64")]
            SimdCapability::Avx2 => unsafe { l2_squared_avx2(a, b) },
            #[cfg(target_arch = "x86_64")]
            SimdCapability::Sse41 => unsafe { l2_squared_sse41(a, b) },
            #[cfg(target_arch = "aarch64")]
            SimdCapability::Neon => unsafe { l2_squared_neon(a, b) },
            _ => l2_squared_scalar(a, b),
        }
    }

    /// Compute dot product between two vectors
    #[inline]
    pub fn dot_product(&self, a: &[f32], b: &[f32]) -> f32 {
        debug_assert_eq!(a.len(), b.len());
        
        match self.capability {
            #[cfg(target_arch = "x86_64")]
            SimdCapability::Avx512 => unsafe { dot_product_avx512(a, b) },
            #[cfg(target_arch = "x86_64")]
            SimdCapability::Avx2 => unsafe { dot_product_avx2(a, b) },
            #[cfg(target_arch = "x86_64")]
            SimdCapability::Sse41 => unsafe { dot_product_sse41(a, b) },
            #[cfg(target_arch = "aarch64")]
            SimdCapability::Neon => unsafe { dot_product_neon(a, b) },
            _ => dot_product_scalar(a, b),
        }
    }

    /// Compute cosine distance (1 - cosine_similarity)
    #[inline]
    pub fn cosine_distance(&self, a: &[f32], b: &[f32]) -> f32 {
        let dot = self.dot_product(a, b);
        let norm_a = self.dot_product(a, a).sqrt();
        let norm_b = self.dot_product(b, b).sqrt();
        
        if norm_a < 1e-10 || norm_b < 1e-10 {
            return 1.0;
        }
        
        1.0 - (dot / (norm_a * norm_b))
    }

    /// Get the SIMD capability being used
    pub fn capability(&self) -> SimdCapability {
        self.capability
    }
}

impl Default for DistanceKernel {
    fn default() -> Self {
        Self::detect()
    }
}

// ============================================================================
// Scalar Implementations (Fallback)
// ============================================================================

/// Scalar L2 squared distance
#[inline]
pub fn l2_squared_scalar(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| {
            let diff = x - y;
            diff * diff
        })
        .sum()
}

/// Scalar dot product
#[inline]
pub fn dot_product_scalar(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

// ============================================================================
// x86_64 AVX2 Implementations
// ============================================================================

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2", enable = "fma")]
#[inline]
pub unsafe fn l2_squared_avx2(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::*;
    
    let n = a.len();
    let mut sum0 = _mm256_setzero_ps();
    let mut sum1 = _mm256_setzero_ps();
    let mut sum2 = _mm256_setzero_ps();
    let mut sum3 = _mm256_setzero_ps();
    
    let chunks = n / 8;
    let chunks4 = chunks / 4;
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    for i in 0..chunks4 {
        let base = i * 32;

        let va0 = _mm256_loadu_ps(a_ptr.add(base));
        let vb0 = _mm256_loadu_ps(b_ptr.add(base));
        let diff0 = _mm256_sub_ps(va0, vb0);
        sum0 = _mm256_fmadd_ps(diff0, diff0, sum0);

        let va1 = _mm256_loadu_ps(a_ptr.add(base + 8));
        let vb1 = _mm256_loadu_ps(b_ptr.add(base + 8));
        let diff1 = _mm256_sub_ps(va1, vb1);
        sum1 = _mm256_fmadd_ps(diff1, diff1, sum1);

        let va2 = _mm256_loadu_ps(a_ptr.add(base + 16));
        let vb2 = _mm256_loadu_ps(b_ptr.add(base + 16));
        let diff2 = _mm256_sub_ps(va2, vb2);
        sum2 = _mm256_fmadd_ps(diff2, diff2, sum2);

        let va3 = _mm256_loadu_ps(a_ptr.add(base + 24));
        let vb3 = _mm256_loadu_ps(b_ptr.add(base + 24));
        let diff3 = _mm256_sub_ps(va3, vb3);
        sum3 = _mm256_fmadd_ps(diff3, diff3, sum3);
    }

    for i in (chunks4 * 4)..chunks {
        let offset = i * 8;
        let va = _mm256_loadu_ps(a_ptr.add(offset));
        let vb = _mm256_loadu_ps(b_ptr.add(offset));
        let diff = _mm256_sub_ps(va, vb);
        sum0 = _mm256_fmadd_ps(diff, diff, sum0);
    }

    let sum01 = _mm256_add_ps(sum0, sum1);
    let sum23 = _mm256_add_ps(sum2, sum3);
    let sum = _mm256_add_ps(sum01, sum23);
    
    // Horizontal sum
    let sum_high = _mm256_extractf128_ps(sum, 1);
    let sum_low = _mm256_castps256_ps128(sum);
    let sum128 = _mm_add_ps(sum_low, sum_high);
    let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    
    let mut result = _mm_cvtss_f32(sum32);
    
    // Handle remainder
    for i in (chunks * 8)..n {
        let diff = *a.get_unchecked(i) - *b.get_unchecked(i);
        result += diff * diff;
    }
    
    result
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2", enable = "fma")]
#[inline]
pub unsafe fn dot_product_avx2(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::*;
    
    let n = a.len();
    let mut sum = _mm256_setzero_ps();
    
    let chunks = n / 8;
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    for i in 0..chunks {
        let offset = i * 8;
        let va = _mm256_loadu_ps(a_ptr.add(offset));
        let vb = _mm256_loadu_ps(b_ptr.add(offset));
        sum = _mm256_fmadd_ps(va, vb, sum);
    }
    
    // Horizontal sum
    let sum_high = _mm256_extractf128_ps(sum, 1);
    let sum_low = _mm256_castps256_ps128(sum);
    let sum128 = _mm_add_ps(sum_low, sum_high);
    let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    
    let mut result = _mm_cvtss_f32(sum32);
    
    // Handle remainder
    for i in (chunks * 8)..n {
        result += *a.get_unchecked(i) * *b.get_unchecked(i);
    }
    
    result
}

// ============================================================================
// x86_64 SSE4.1 Implementations
// ============================================================================

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.1")]
#[inline]
pub unsafe fn l2_squared_sse41(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::*;
    
    let n = a.len();
    let mut sum = _mm_setzero_ps();
    
    let chunks = n / 4;
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    for i in 0..chunks {
        let offset = i * 4;
        let va = _mm_loadu_ps(a_ptr.add(offset));
        let vb = _mm_loadu_ps(b_ptr.add(offset));
        let diff = _mm_sub_ps(va, vb);
        let sq = _mm_mul_ps(diff, diff);
        sum = _mm_add_ps(sum, sq);
    }
    
    // Horizontal sum
    let sum64 = _mm_add_ps(sum, _mm_movehl_ps(sum, sum));
    let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    
    let mut result = _mm_cvtss_f32(sum32);
    
    // Handle remainder
    for i in (chunks * 4)..n {
        let diff = *a.get_unchecked(i) - *b.get_unchecked(i);
        result += diff * diff;
    }
    
    result
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.1")]
#[inline]
pub unsafe fn dot_product_sse41(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::*;
    
    let n = a.len();
    let mut sum = _mm_setzero_ps();
    
    let chunks = n / 4;
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    for i in 0..chunks {
        let offset = i * 4;
        let va = _mm_loadu_ps(a_ptr.add(offset));
        let vb = _mm_loadu_ps(b_ptr.add(offset));
        let prod = _mm_mul_ps(va, vb);
        sum = _mm_add_ps(sum, prod);
    }
    
    // Horizontal sum
    let sum64 = _mm_add_ps(sum, _mm_movehl_ps(sum, sum));
    let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    
    let mut result = _mm_cvtss_f32(sum32);
    
    // Handle remainder
    for i in (chunks * 4)..n {
        result += *a.get_unchecked(i) * *b.get_unchecked(i);
    }
    
    result
}

// ============================================================================
// x86_64 AVX-512 Implementations
// ============================================================================

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
#[inline]
pub unsafe fn l2_squared_avx512(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::*;
    
    let n = a.len();
    let mut sum = _mm512_setzero_ps();
    
    let chunks = n / 16;
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    for i in 0..chunks {
        let offset = i * 16;
        let va = _mm512_loadu_ps(a_ptr.add(offset));
        let vb = _mm512_loadu_ps(b_ptr.add(offset));
        let diff = _mm512_sub_ps(va, vb);
        sum = _mm512_fmadd_ps(diff, diff, sum);
    }
    
    // Reduce to scalar
    let mut result = _mm512_reduce_add_ps(sum);
    
    // Handle remainder
    for i in (chunks * 16)..n {
        let diff = *a.get_unchecked(i) - *b.get_unchecked(i);
        result += diff * diff;
    }
    
    result
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
#[inline]
pub unsafe fn dot_product_avx512(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::*;
    
    let n = a.len();
    let mut sum = _mm512_setzero_ps();
    
    let chunks = n / 16;
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    for i in 0..chunks {
        let offset = i * 16;
        let va = _mm512_loadu_ps(a_ptr.add(offset));
        let vb = _mm512_loadu_ps(b_ptr.add(offset));
        sum = _mm512_fmadd_ps(va, vb, sum);
    }
    
    // Reduce to scalar
    let mut result = _mm512_reduce_add_ps(sum);
    
    // Handle remainder
    for i in (chunks * 16)..n {
        result += *a.get_unchecked(i) * *b.get_unchecked(i);
    }
    
    result
}

// ============================================================================
// aarch64 NEON Implementations
// ============================================================================

#[cfg(target_arch = "aarch64")]
#[inline]
pub unsafe fn l2_squared_neon(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::*;
    
    unsafe {
        let n = a.len();
        let mut sum = vdupq_n_f32(0.0);
        
        let chunks = n / 4;
        let a_ptr = a.as_ptr();
        let b_ptr = b.as_ptr();
        
        for i in 0..chunks {
            let offset = i * 4;
            let va = vld1q_f32(a_ptr.add(offset));
            let vb = vld1q_f32(b_ptr.add(offset));
            let diff = vsubq_f32(va, vb);
            sum = vfmaq_f32(sum, diff, diff);
        }
        
        // Horizontal sum
        let mut result = vaddvq_f32(sum);
        
        // Handle remainder
        for i in (chunks * 4)..n {
            let diff = *a.get_unchecked(i) - *b.get_unchecked(i);
            result += diff * diff;
        }
        
        result
    }
}

#[cfg(target_arch = "aarch64")]
#[inline]
pub unsafe fn dot_product_neon(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::*;
    
    unsafe {
        let n = a.len();
        let mut sum = vdupq_n_f32(0.0);
        
        let chunks = n / 4;
        let a_ptr = a.as_ptr();
        let b_ptr = b.as_ptr();
        
        for i in 0..chunks {
            let offset = i * 4;
            let va = vld1q_f32(a_ptr.add(offset));
            let vb = vld1q_f32(b_ptr.add(offset));
            sum = vfmaq_f32(sum, va, vb);
        }
        
        // Horizontal sum
        let mut result = vaddvq_f32(sum);
        
        // Handle remainder
        for i in (chunks * 4)..n {
            result += *a.get_unchecked(i) * *b.get_unchecked(i);
        }
        
        result
    }
}

/// Task #7: High-ILP NEON fused cosine distance (computes dot, norm_a, norm_b in one pass)
/// Uses 6 accumulators (2 per metric) to saturate FMLA throughput on Apple Silicon
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
#[inline]
pub unsafe fn cosine_distance_neon_fused(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::*;
    
    unsafe {
        let n = a.len();
        
        // 6 accumulators for ILP: 2 each for dot, norm_a, norm_b
        let mut dot0 = vdupq_n_f32(0.0);
        let mut dot1 = vdupq_n_f32(0.0);
        let mut norm_a0 = vdupq_n_f32(0.0);
        let mut norm_a1 = vdupq_n_f32(0.0);
        let mut norm_b0 = vdupq_n_f32(0.0);
        let mut norm_b1 = vdupq_n_f32(0.0);
        
        let a_ptr = a.as_ptr();
        let b_ptr = b.as_ptr();
        
        // Process 8 floats per iteration (2 × 4-wide NEON)
        let mut i = 0;
        while i + 8 <= n {
            let va0 = vld1q_f32(a_ptr.add(i));
            let vb0 = vld1q_f32(b_ptr.add(i));
            let va1 = vld1q_f32(a_ptr.add(i + 4));
            let vb1 = vld1q_f32(b_ptr.add(i + 4));
            
            // All 6 FMLAs are independent - CPU can issue in parallel
            dot0 = vfmaq_f32(dot0, va0, vb0);
            dot1 = vfmaq_f32(dot1, va1, vb1);
            norm_a0 = vfmaq_f32(norm_a0, va0, va0);
            norm_a1 = vfmaq_f32(norm_a1, va1, va1);
            norm_b0 = vfmaq_f32(norm_b0, vb0, vb0);
            norm_b1 = vfmaq_f32(norm_b1, vb1, vb1);
            
            i += 8;
        }
        
        // Process remaining 4-element chunks
        while i + 4 <= n {
            let va = vld1q_f32(a_ptr.add(i));
            let vb = vld1q_f32(b_ptr.add(i));
            dot0 = vfmaq_f32(dot0, va, vb);
            norm_a0 = vfmaq_f32(norm_a0, va, va);
            norm_b0 = vfmaq_f32(norm_b0, vb, vb);
            i += 4;
        }
        
        // Merge accumulators
        let dot_sum = vaddq_f32(dot0, dot1);
        let norm_a_sum = vaddq_f32(norm_a0, norm_a1);
        let norm_b_sum = vaddq_f32(norm_b0, norm_b1);
        
        let mut dot = vaddvq_f32(dot_sum);
        let mut norm_a_sq = vaddvq_f32(norm_a_sum);
        let mut norm_b_sq = vaddvq_f32(norm_b_sum);
        
        // Handle remainder
        for j in i..n {
            let av = *a.get_unchecked(j);
            let bv = *b.get_unchecked(j);
            dot += av * bv;
            norm_a_sq += av * av;
            norm_b_sq += bv * bv;
        }
        
        let norm_a = norm_a_sq.sqrt();
        let norm_b = norm_b_sq.sqrt();
        
        if norm_a < 1e-10 || norm_b < 1e-10 {
            1.0
        } else {
            1.0 - (dot / (norm_a * norm_b))
        }
    }
}

// ============================================================================
// Convenience Functions
// ============================================================================

/// Compute L2 squared distance with auto-detected SIMD
#[inline]
pub fn l2_squared(a: &[f32], b: &[f32]) -> f32 {
    DistanceKernel::detect().l2_squared(a, b)
}

/// Compute dot product with auto-detected SIMD
#[inline]
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
    DistanceKernel::detect().dot_product(a, b)
}

/// Compute cosine distance with auto-detected SIMD
#[inline]
pub fn cosine_distance(a: &[f32], b: &[f32]) -> f32 {
    DistanceKernel::detect().cosine_distance(a, b)
}

// ============================================================================
// ULTRA-FAST STATIC DISPATCH (Zero overhead - used by search_fast)
// ============================================================================

/// Static cached kernel for zero-overhead dispatch
static CACHED_KERNEL: OnceLock<DistanceKernel> = OnceLock::new();

fn get_kernel() -> &'static DistanceKernel {
    CACHED_KERNEL.get_or_init(DistanceKernel::detect)
}

/// Ultra-fast L2 distance (static dispatch, no heap allocation)
#[inline(always)]
pub fn l2_distance_fast(a: &[f32], b: &[f32]) -> f32 {
    get_kernel().l2_squared(a, b).sqrt()
}

/// Ultra-fast dot product (static dispatch, no heap allocation)
#[inline(always)]
pub fn dot_product_fast(a: &[f32], b: &[f32]) -> f32 {
    get_kernel().dot_product(a, b)
}

/// Ultra-fast cosine distance (static dispatch, no heap allocation)  
/// 
/// For normalized vectors (unit length), this is just 1 - dot_product.
/// For non-normalized vectors, computes full cosine distance.
#[inline(always)]
pub fn cosine_distance_fast(a: &[f32], b: &[f32]) -> f32 {
    // Check if we're on aarch64 and can use fused implementation
    #[cfg(target_arch = "aarch64")]
    {
        return unsafe { cosine_distance_neon_fused(a, b) };
    }
    
    #[cfg(not(target_arch = "aarch64"))]
    {
        get_kernel().cosine_distance(a, b)
    }
}

/// Threshold-aware L2 squared distance with early abort
/// Returns the actual distance if it's <= threshold, or a value > threshold if it exceeds
/// This enables branch-and-bound optimization for RNG selection
#[inline]
pub fn l2_squared_threshold(a: &[f32], b: &[f32], threshold_squared: f32) -> f32 {
    if a.len() != b.len() {
        return f32::INFINITY;
    }
    
    // Try AVX2 path if available and worthwhile
    #[cfg(target_arch = "x86_64")]
    if is_x86_feature_detected!("avx2") && a.len() >= 8 {
        return unsafe { l2_squared_threshold_avx2(a, b, threshold_squared) };
    }

    // Scalar fallback with early abort
    l2_squared_threshold_scalar(a, b, threshold_squared)
}

/// Scalar implementation with early abort
#[inline]
fn l2_squared_threshold_scalar(a: &[f32], b: &[f32], threshold_squared: f32) -> f32 {
    let mut sum = 0.0f32;
    
    for (x, y) in a.iter().zip(b.iter()) {
        let diff = x - y;
        sum += diff * diff;
        
        // Early abort if we've exceeded threshold
        if sum > threshold_squared {
            return sum;
        }
    }
    
    sum
}

/// AVX2 implementation with early abort
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn l2_squared_threshold_avx2(a: &[f32], b: &[f32], threshold_squared: f32) -> f32 {
    use std::arch::x86_64::*;

    let n = a.len();
    let chunks = n / 8;
    let _remainder = n % 8;
    
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();
    
    let mut sum = _mm256_setzero_ps();
    let _threshold_vec = _mm256_set1_ps(threshold_squared);
    
    // Process 8 elements at a time with early abort checking every few chunks
    let check_interval = 4; // Check threshold every 32 elements
    
    for chunk_group in (0..chunks).step_by(check_interval) {
        let end_chunk = std::cmp::min(chunk_group + check_interval, chunks);
        
        for chunk in chunk_group..end_chunk {
            let offset = chunk * 8;
            unsafe {
                let va = _mm256_loadu_ps(a_ptr.add(offset));
                let vb = _mm256_loadu_ps(b_ptr.add(offset));
                let diff = _mm256_sub_ps(va, vb);
                // Use mul + add instead of fmadd to avoid target_feature requirement
                let sq = _mm256_mul_ps(diff, diff);
                sum = _mm256_add_ps(sum, sq);
            }
        }
        
        // Check if we've exceeded threshold
        let sum_scalar = unsafe {
            let sum_high = _mm256_extractf128_ps(sum, 1);
            let sum_low = _mm256_castps256_ps128(sum);
            let sum128 = _mm_add_ps(sum_low, sum_high);
            let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
            let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
            _mm_cvtss_f32(sum32)
        };
        
        if sum_scalar > threshold_squared {
            // Add remainder and return (we know we've exceeded threshold)
            for i in (end_chunk * 8)..n {
                let diff = a[i] - b[i];
                sum_scalar + diff * diff;
            }
            return sum_scalar;
        }
    }
    
    // Horizontal sum for final result
    let mut result = unsafe {
        let sum_high = _mm256_extractf128_ps(sum, 1);
        let sum_low = _mm256_castps256_ps128(sum);
        let sum128 = _mm_add_ps(sum_low, sum_high);
        let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
        let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
        _mm_cvtss_f32(sum32)
    };
    
    // Handle remainder
    for i in (chunks * 8)..n {
        let diff = a[i] - b[i];
        result += diff * diff;
        
        // Early abort check for remainder
        if result > threshold_squared {
            return result;
        }
    }
    
    result
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    fn random_vec(n: usize, seed: u64) -> Vec<f32> {
        (0..n)
            .map(|i| ((i as u64 + seed) % 1000) as f32 / 1000.0 - 0.5)
            .collect()
    }

    #[test]
    fn test_l2_squared_scalar() {
        let a = vec![1.0, 2.0, 3.0, 4.0];
        let b = vec![5.0, 6.0, 7.0, 8.0];
        
        let result = l2_squared_scalar(&a, &b);
        // (5-1)^2 + (6-2)^2 + (7-3)^2 + (8-4)^2 = 16 + 16 + 16 + 16 = 64
        assert!((result - 64.0).abs() < 0.01);
    }

    #[test]
    fn test_dot_product_scalar() {
        let a = vec![1.0, 2.0, 3.0, 4.0];
        let b = vec![5.0, 6.0, 7.0, 8.0];
        
        let result = dot_product_scalar(&a, &b);
        // 1*5 + 2*6 + 3*7 + 4*8 = 5 + 12 + 21 + 32 = 70
        assert!((result - 70.0).abs() < 0.01);
    }

    #[test]
    fn test_kernel_detection() {
        let kernel = DistanceKernel::detect();
        let cap = kernel.capability();
        
        // Should detect something
        #[cfg(target_arch = "x86_64")]
        assert!(matches!(cap, SimdCapability::Scalar | SimdCapability::Sse41 | SimdCapability::Avx2 | SimdCapability::Avx512));
        
        #[cfg(target_arch = "aarch64")]
        assert_eq!(cap, SimdCapability::Neon);
    }

    #[test]
    fn test_l2_squared_consistency() {
        let a = random_vec(768, 42);
        let b = random_vec(768, 123);
        
        let scalar_result = l2_squared_scalar(&a, &b);
        let kernel_result = l2_squared(&a, &b);
        
        // Results should match within floating point tolerance
        let rel_error = (scalar_result - kernel_result).abs() / scalar_result.max(1e-10);
        assert!(rel_error < 1e-5, "Relative error: {}", rel_error);
    }

    #[test]
    fn test_dot_product_consistency() {
        let a = random_vec(768, 42);
        let b = random_vec(768, 123);
        
        let scalar_result = dot_product_scalar(&a, &b);
        let kernel_result = dot_product(&a, &b);
        
        // Results should match within floating point tolerance
        let rel_error = (scalar_result - kernel_result).abs() / scalar_result.abs().max(1e-10);
        assert!(rel_error < 1e-5, "Relative error: {}", rel_error);
    }

    #[test]
    fn test_cosine_distance() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        
        let result = cosine_distance(&a, &b);
        assert!(result.abs() < 0.01); // Same direction = distance 0
        
        let c = vec![0.0, 1.0, 0.0];
        let result2 = cosine_distance(&a, &c);
        assert!((result2 - 1.0).abs() < 0.01); // Orthogonal = distance 1
    }
}
