//! Distance functions for vector similarity
//!
//! Provides efficient implementations of common distance metrics:
//! - Cosine similarity/distance
//! - Euclidean (L2) distance
//! - Dot product (inner product)
//!
//! P1-2: Runtime SIMD detection for AVX2, AVX512, and NEON

use crate::types::DistanceMetric;

// ============================================================================
// SIMD Feature Detection (P1-2)
// ============================================================================

/// Detected SIMD capabilities at runtime
#[derive(Debug, Clone, Copy)]
pub struct SimdCapabilities {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub avx2: bool,
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub avx512f: bool,
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub fma: bool,
    #[cfg(target_arch = "aarch64")]
    pub neon: bool,
}

impl SimdCapabilities {
    /// Detect SIMD capabilities at runtime
    pub fn detect() -> Self {
        Self {
            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
            avx2: std::arch::is_x86_feature_detected!("avx2"),
            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
            avx512f: std::arch::is_x86_feature_detected!("avx512f"),
            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
            fma: std::arch::is_x86_feature_detected!("fma"),
            #[cfg(target_arch = "aarch64")]
            neon: true, // NEON is always available on aarch64
        }
    }
    
    /// Get a description of available SIMD features
    pub fn description(&self) -> String {
        let mut features = Vec::new();
        
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        {
            if self.avx512f {
                features.push("AVX-512F");
            }
            if self.avx2 {
                features.push("AVX2");
            }
            if self.fma {
                features.push("FMA");
            }
        }
        
        #[cfg(target_arch = "aarch64")]
        {
            if self.neon {
                features.push("NEON");
            }
        }
        
        if features.is_empty() {
            "Scalar".to_string()
        } else {
            features.join(", ")
        }
    }
}

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

/// Distance calculation trait
pub trait Distance: Send + Sync {
    /// Calculate distance between two vectors
    /// Lower values mean more similar for distance metrics
    /// Higher values mean more similar for similarity metrics (converted internally)
    fn calculate(&self, a: &[f32], b: &[f32]) -> f32;
    
    /// Get the metric type
    fn metric(&self) -> DistanceMetric;
}

/// Distance calculator with SIMD-optimized implementations
/// 
/// Automatically detects and uses the best available SIMD instructions
#[derive(Debug, Clone)]
pub struct DistanceCalculator {
    metric: DistanceMetric,
    simd: SimdCapabilities,
}

impl DistanceCalculator {
    /// Create a new distance calculator with automatic SIMD detection
    pub fn new(metric: DistanceMetric) -> Self {
        Self { 
            metric,
            simd: SimdCapabilities::detect(),
        }
    }
    
    /// Create with cosine distance
    pub fn cosine() -> Self {
        Self::new(DistanceMetric::Cosine)
    }
    
    /// Create with Euclidean distance
    pub fn euclidean() -> Self {
        Self::new(DistanceMetric::Euclidean)
    }
    
    /// Create with dot product
    pub fn dot_product() -> Self {
        Self::new(DistanceMetric::DotProduct)
    }
    
    /// Get SIMD capabilities
    pub fn simd_capabilities(&self) -> &SimdCapabilities {
        &self.simd
    }
    
    /// Calculate dot product using the best available SIMD
    #[inline]
    fn dot_product_simd(&self, a: &[f32], b: &[f32]) -> f32 {
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        {
            if self.simd.avx2 && self.simd.fma {
                return unsafe { simd_x86::dot_product_avx2_fma(a, b) };
            }
            if self.simd.avx2 {
                return unsafe { simd_x86::dot_product_avx2(a, b) };
            }
        }
        
        #[cfg(target_arch = "aarch64")]
        {
            if self.simd.neon {
                return unsafe { simd_neon::dot_product_neon(a, b) };
            }
        }
        
        // Fallback to scalar
        dot_product_scalar(a, b)
    }
    
    /// Calculate squared Euclidean distance using the best available SIMD
    #[inline]
    fn euclidean_squared_simd(&self, a: &[f32], b: &[f32]) -> f32 {
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        {
            if self.simd.avx2 && self.simd.fma {
                return unsafe { simd_x86::euclidean_squared_avx2_fma(a, b) };
            }
            if self.simd.avx2 {
                return unsafe { simd_x86::euclidean_squared_avx2(a, b) };
            }
        }
        
        #[cfg(target_arch = "aarch64")]
        {
            if self.simd.neon {
                return unsafe { simd_neon::euclidean_squared_neon(a, b) };
            }
        }
        
        // Fallback to scalar
        euclidean_distance_squared_scalar(a, b)
    }
    
    /// Calculate dot product and norms using SIMD
    #[inline]
    fn dot_and_norms_simd(&self, a: &[f32], b: &[f32]) -> (f32, f32, f32) {
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        {
            if self.simd.avx2 {
                return unsafe { simd_x86::dot_and_norms_avx2(a, b) };
            }
        }
        
        #[cfg(target_arch = "aarch64")]
        {
            if self.simd.neon {
                return unsafe { simd_neon::dot_and_norms_neon(a, b) };
            }
        }
        
        // Fallback to scalar
        dot_and_norms_scalar(a, b)
    }
}

impl Distance for DistanceCalculator {
    fn calculate(&self, a: &[f32], b: &[f32]) -> f32 {
        debug_assert_eq!(a.len(), b.len(), "Vector dimensions must match");
        
        match self.metric {
            DistanceMetric::Cosine => {
                let (dot, norm_a, norm_b) = self.dot_and_norms_simd(a, b);
                let denom = (norm_a * norm_b).sqrt();
                if denom < f32::EPSILON {
                    return 1.0;
                }
                1.0 - (dot / denom)
            }
            DistanceMetric::Euclidean => self.euclidean_squared_simd(a, b).sqrt(),
            DistanceMetric::DotProduct => -self.dot_product_simd(a, b),
        }
    }
    
    fn metric(&self) -> DistanceMetric {
        self.metric
    }
}

// ============================================================================
// Scalar Implementations
// ============================================================================

/// Calculate cosine distance (1 - cosine_similarity) - scalar version
/// Returns 0 for identical vectors, 2 for opposite vectors
#[inline]
pub fn cosine_distance(a: &[f32], b: &[f32]) -> f32 {
    let (dot, norm_a, norm_b) = dot_and_norms_scalar(a, b);
    
    let denom = (norm_a * norm_b).sqrt();
    if denom < f32::EPSILON {
        return 1.0; // Handle zero vectors
    }
    
    1.0 - (dot / denom)
}

/// Calculate cosine similarity (-1 to 1, higher is more similar)
#[inline]
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    1.0 - cosine_distance(a, b)
}

/// Calculate Euclidean (L2) distance
#[inline]
pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
    euclidean_distance_squared_scalar(a, b).sqrt()
}

/// Calculate squared Euclidean distance (faster, avoids sqrt) - scalar version
#[inline]
pub fn euclidean_distance_squared(a: &[f32], b: &[f32]) -> f32 {
    euclidean_distance_squared_scalar(a, b)
}

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

/// Calculate dot product (inner product)
#[inline]
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
    dot_product_scalar(a, b)
}

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

/// Negative dot product for use as distance (higher similarity = lower distance)
#[inline]
pub fn negative_dot_product(a: &[f32], b: &[f32]) -> f32 {
    -dot_product(a, b)
}

/// Calculate dot product and squared norms in a single pass - scalar version
#[inline]
fn dot_and_norms_scalar(a: &[f32], b: &[f32]) -> (f32, f32, f32) {
    let mut dot = 0.0f32;
    let mut norm_a = 0.0f32;
    let mut norm_b = 0.0f32;
    
    for (x, y) in a.iter().zip(b.iter()) {
        dot += x * y;
        norm_a += x * x;
        norm_b += y * y;
    }
    
    (dot, norm_a, norm_b)
}

/// Calculate L2 norm of a vector
#[inline]
pub fn l2_norm(v: &[f32]) -> f32 {
    v.iter().map(|x| x * x).sum::<f32>().sqrt()
}

/// Normalize a vector to unit length (in place)
#[inline]
pub fn normalize(v: &mut [f32]) {
    let norm = l2_norm(v);
    if norm > f32::EPSILON {
        for x in v.iter_mut() {
            *x /= norm;
        }
    }
}

/// Normalize a vector to unit length (returns new vector)
#[inline]
pub fn normalized(v: &[f32]) -> Vec<f32> {
    let norm = l2_norm(v);
    if norm < f32::EPSILON {
        return v.to_vec();
    }
    v.iter().map(|x| x / norm).collect()
}

// ============================================================================
// SIMD-optimized versions (P1-2)
// ============================================================================

/// x86/x86_64 SIMD implementations
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
mod simd_x86 {
    #[cfg(target_arch = "x86")]
    use std::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    
    /// AVX2 dot product without FMA
    #[target_feature(enable = "avx2")]
    pub unsafe fn dot_product_avx2(a: &[f32], b: &[f32]) -> f32 {
        let n = a.len();
        let chunks = n / 8;
        
        let mut sum = _mm256_setzero_ps();
        
        for i in 0..chunks {
            let offset = i * 8;
            let va = _mm256_loadu_ps(a.as_ptr().add(offset));
            let vb = _mm256_loadu_ps(b.as_ptr().add(offset));
            let prod = _mm256_mul_ps(va, vb);
            sum = _mm256_add_ps(sum, prod);
        }
        
        // Horizontal sum
        let mut result = [0.0f32; 8];
        _mm256_storeu_ps(result.as_mut_ptr(), sum);
        let mut total: f32 = result.iter().sum();
        
        // Handle remainder
        for i in (chunks * 8)..n {
            total += a[i] * b[i];
        }
        
        total
    }
    
    /// AVX2 + FMA dot product (faster with fused multiply-add)
    #[target_feature(enable = "avx2", enable = "fma")]
    pub unsafe fn dot_product_avx2_fma(a: &[f32], b: &[f32]) -> f32 {
        let n = a.len();
        let chunks = n / 8;
        
        let mut sum = _mm256_setzero_ps();
        
        for i in 0..chunks {
            let offset = i * 8;
            let va = _mm256_loadu_ps(a.as_ptr().add(offset));
            let vb = _mm256_loadu_ps(b.as_ptr().add(offset));
            sum = _mm256_fmadd_ps(va, vb, sum);
        }
        
        // Horizontal sum
        let mut result = [0.0f32; 8];
        _mm256_storeu_ps(result.as_mut_ptr(), sum);
        let mut total: f32 = result.iter().sum();
        
        // Handle remainder
        for i in (chunks * 8)..n {
            total += a[i] * b[i];
        }
        
        total
    }
    
    /// AVX2 squared Euclidean distance
    #[target_feature(enable = "avx2")]
    pub unsafe fn euclidean_squared_avx2(a: &[f32], b: &[f32]) -> f32 {
        let n = a.len();
        let chunks = n / 8;
        
        let mut sum = _mm256_setzero_ps();
        
        for i in 0..chunks {
            let offset = i * 8;
            let va = _mm256_loadu_ps(a.as_ptr().add(offset));
            let vb = _mm256_loadu_ps(b.as_ptr().add(offset));
            let diff = _mm256_sub_ps(va, vb);
            let sq = _mm256_mul_ps(diff, diff);
            sum = _mm256_add_ps(sum, sq);
        }
        
        // Horizontal sum
        let mut result = [0.0f32; 8];
        _mm256_storeu_ps(result.as_mut_ptr(), sum);
        let mut total: f32 = result.iter().sum();
        
        // Handle remainder
        for i in (chunks * 8)..n {
            let diff = a[i] - b[i];
            total += diff * diff;
        }
        
        total
    }
    
    /// AVX2 + FMA squared Euclidean distance
    #[target_feature(enable = "avx2", enable = "fma")]
    pub unsafe fn euclidean_squared_avx2_fma(a: &[f32], b: &[f32]) -> f32 {
        let n = a.len();
        let chunks = n / 8;
        
        let mut sum = _mm256_setzero_ps();
        
        for i in 0..chunks {
            let offset = i * 8;
            let va = _mm256_loadu_ps(a.as_ptr().add(offset));
            let vb = _mm256_loadu_ps(b.as_ptr().add(offset));
            let diff = _mm256_sub_ps(va, vb);
            sum = _mm256_fmadd_ps(diff, diff, sum);
        }
        
        // Horizontal sum
        let mut result = [0.0f32; 8];
        _mm256_storeu_ps(result.as_mut_ptr(), sum);
        let mut total: f32 = result.iter().sum();
        
        // Handle remainder
        for i in (chunks * 8)..n {
            let diff = a[i] - b[i];
            total += diff * diff;
        }
        
        total
    }
    
    /// AVX2 dot product and norms in single pass
    #[target_feature(enable = "avx2")]
    pub unsafe fn dot_and_norms_avx2(a: &[f32], b: &[f32]) -> (f32, f32, f32) {
        let n = a.len();
        let chunks = n / 8;
        
        let mut dot_sum = _mm256_setzero_ps();
        let mut norm_a_sum = _mm256_setzero_ps();
        let mut norm_b_sum = _mm256_setzero_ps();
        
        for i in 0..chunks {
            let offset = i * 8;
            let va = _mm256_loadu_ps(a.as_ptr().add(offset));
            let vb = _mm256_loadu_ps(b.as_ptr().add(offset));
            
            // dot += a * b
            let prod = _mm256_mul_ps(va, vb);
            dot_sum = _mm256_add_ps(dot_sum, prod);
            
            // norm_a += a * a
            let sq_a = _mm256_mul_ps(va, va);
            norm_a_sum = _mm256_add_ps(norm_a_sum, sq_a);
            
            // norm_b += b * b
            let sq_b = _mm256_mul_ps(vb, vb);
            norm_b_sum = _mm256_add_ps(norm_b_sum, sq_b);
        }
        
        // Horizontal sums
        let mut dot_arr = [0.0f32; 8];
        let mut norm_a_arr = [0.0f32; 8];
        let mut norm_b_arr = [0.0f32; 8];
        
        _mm256_storeu_ps(dot_arr.as_mut_ptr(), dot_sum);
        _mm256_storeu_ps(norm_a_arr.as_mut_ptr(), norm_a_sum);
        _mm256_storeu_ps(norm_b_arr.as_mut_ptr(), norm_b_sum);
        
        let mut dot: f32 = dot_arr.iter().sum();
        let mut norm_a: f32 = norm_a_arr.iter().sum();
        let mut norm_b: f32 = norm_b_arr.iter().sum();
        
        // Handle remainder
        for i in (chunks * 8)..n {
            dot += a[i] * b[i];
            norm_a += a[i] * a[i];
            norm_b += b[i] * b[i];
        }
        
        (dot, norm_a, norm_b)
    }
}

/// ARM NEON SIMD implementations
#[cfg(target_arch = "aarch64")]
mod simd_neon {
    use std::arch::aarch64::*;
    
    /// NEON dot product
    #[target_feature(enable = "neon")]
    pub unsafe fn dot_product_neon(a: &[f32], b: &[f32]) -> f32 {
        let n = a.len();
        let chunks = n / 4;
        
        let mut sum = vdupq_n_f32(0.0);
        
        for i in 0..chunks {
            let offset = i * 4;
            let va = vld1q_f32(a.as_ptr().add(offset));
            let vb = vld1q_f32(b.as_ptr().add(offset));
            sum = vfmaq_f32(sum, va, vb);
        }
        
        // Horizontal sum
        let mut total = vaddvq_f32(sum);
        
        // Handle remainder
        for i in (chunks * 4)..n {
            total += a[i] * b[i];
        }
        
        total
    }
    
    /// NEON squared Euclidean distance
    #[target_feature(enable = "neon")]
    pub unsafe fn euclidean_squared_neon(a: &[f32], b: &[f32]) -> f32 {
        let n = a.len();
        let chunks = n / 4;
        
        let mut sum = vdupq_n_f32(0.0);
        
        for i in 0..chunks {
            let offset = i * 4;
            let va = vld1q_f32(a.as_ptr().add(offset));
            let vb = vld1q_f32(b.as_ptr().add(offset));
            let diff = vsubq_f32(va, vb);
            sum = vfmaq_f32(sum, diff, diff);
        }
        
        // Horizontal sum
        let mut total = vaddvq_f32(sum);
        
        // Handle remainder
        for i in (chunks * 4)..n {
            let diff = a[i] - b[i];
            total += diff * diff;
        }
        
        total
    }
    
    /// NEON dot product and norms
    #[target_feature(enable = "neon")]
    pub unsafe fn dot_and_norms_neon(a: &[f32], b: &[f32]) -> (f32, f32, f32) {
        let n = a.len();
        let chunks = n / 4;
        
        let mut dot_sum = vdupq_n_f32(0.0);
        let mut norm_a_sum = vdupq_n_f32(0.0);
        let mut norm_b_sum = vdupq_n_f32(0.0);
        
        for i in 0..chunks {
            let offset = i * 4;
            let va = vld1q_f32(a.as_ptr().add(offset));
            let vb = vld1q_f32(b.as_ptr().add(offset));
            
            dot_sum = vfmaq_f32(dot_sum, va, vb);
            norm_a_sum = vfmaq_f32(norm_a_sum, va, va);
            norm_b_sum = vfmaq_f32(norm_b_sum, vb, vb);
        }
        
        // Horizontal sums
        let mut dot = vaddvq_f32(dot_sum);
        let mut norm_a = vaddvq_f32(norm_a_sum);
        let mut norm_b = vaddvq_f32(norm_b_sum);
        
        // Handle remainder
        for i in (chunks * 4)..n {
            dot += a[i] * b[i];
            norm_a += a[i] * a[i];
            norm_b += b[i] * b[i];
        }
        
        (dot, norm_a, norm_b)
    }
}

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

    fn approx_eq(a: f32, b: f32, epsilon: f32) -> bool {
        (a - b).abs() < epsilon
    }

    #[test]
    fn test_cosine_distance() {
        // Identical vectors should have distance 0
        let a = vec![1.0, 0.0, 0.0];
        assert!(approx_eq(cosine_distance(&a, &a), 0.0, 1e-6));
        
        // Orthogonal vectors should have distance 1
        let b = vec![0.0, 1.0, 0.0];
        assert!(approx_eq(cosine_distance(&a, &b), 1.0, 1e-6));
        
        // Opposite vectors should have distance 2
        let c = vec![-1.0, 0.0, 0.0];
        assert!(approx_eq(cosine_distance(&a, &c), 2.0, 1e-6));
    }

    #[test]
    fn test_euclidean_distance() {
        let a = vec![0.0, 0.0, 0.0];
        let b = vec![1.0, 0.0, 0.0];
        assert!(approx_eq(euclidean_distance(&a, &b), 1.0, 1e-6));
        
        let c = vec![3.0, 4.0, 0.0];
        assert!(approx_eq(euclidean_distance(&a, &c), 5.0, 1e-6));
    }

    #[test]
    fn test_dot_product() {
        let a = vec![1.0, 2.0, 3.0];
        let b = vec![4.0, 5.0, 6.0];
        assert!(approx_eq(dot_product(&a, &b), 32.0, 1e-6));
    }

    #[test]
    fn test_normalize() {
        let v = vec![3.0, 4.0];
        let normalized_v = normalized(&v);
        assert!(approx_eq(l2_norm(&normalized_v), 1.0, 1e-6));
        assert!(approx_eq(normalized_v[0], 0.6, 1e-6));
        assert!(approx_eq(normalized_v[1], 0.8, 1e-6));
    }

    #[test]
    fn test_distance_calculator() {
        let calc = DistanceCalculator::cosine();
        let a = vec![1.0, 0.0];
        let b = vec![0.0, 1.0];
        
        let dist = calc.calculate(&a, &b);
        assert!(approx_eq(dist, 1.0, 1e-6));
        assert_eq!(calc.metric(), DistanceMetric::Cosine);
    }
    
    #[test]
    fn test_zero_vector_handling() {
        let zero = vec![0.0, 0.0, 0.0];
        let nonzero = vec![1.0, 2.0, 3.0];
        
        // Should not panic, return reasonable value
        let dist = cosine_distance(&zero, &nonzero);
        assert!(dist.is_finite());
    }
    
    #[test]
    fn test_simd_capabilities_detection() {
        let caps = SimdCapabilities::detect();
        let desc = caps.description();
        
        // Should return some description
        assert!(!desc.is_empty());
        println!("Detected SIMD capabilities: {}", desc);
    }
    
    #[test]
    fn test_simd_distance_calculator() {
        let calc = DistanceCalculator::new(DistanceMetric::Cosine);
        
        // Test with various vector sizes including non-multiple of SIMD width
        for size in [3, 8, 16, 17, 32, 33, 64, 100, 128, 256, 1536].iter() {
            let a: Vec<f32> = (0..*size).map(|i| (i as f32) * 0.01).collect();
            let b: Vec<f32> = (0..*size).map(|i| (i as f32) * 0.02 + 0.5).collect();
            
            // Compare SIMD result with scalar result
            let simd_result = calc.calculate(&a, &b);
            let scalar_result = cosine_distance(&a, &b);
            
            assert!(
                approx_eq(simd_result, scalar_result, 1e-5),
                "SIMD and scalar results differ for size {}: {} vs {}",
                size, simd_result, scalar_result
            );
        }
    }
    
    #[test]
    fn test_simd_euclidean() {
        let calc = DistanceCalculator::new(DistanceMetric::Euclidean);
        
        let a: Vec<f32> = (0..128).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..128).map(|i| (i as f32) + 1.0).collect();
        
        let simd_result = calc.calculate(&a, &b);
        let scalar_result = euclidean_distance(&a, &b);
        
        assert!(
            approx_eq(simd_result, scalar_result, 1e-4),
            "SIMD and scalar Euclidean differ: {} vs {}",
            simd_result, scalar_result
        );
    }
    
    #[test]
    fn test_simd_dot_product() {
        let calc = DistanceCalculator::new(DistanceMetric::DotProduct);
        
        let a: Vec<f32> = (0..128).map(|i| i as f32 * 0.1).collect();
        let b: Vec<f32> = (0..128).map(|i| i as f32 * 0.2).collect();
        
        let simd_result = calc.calculate(&a, &b);
        let scalar_result = negative_dot_product(&a, &b);
        
        assert!(
            approx_eq(simd_result, scalar_result, 1e-3),
            "SIMD and scalar dot product differ: {} vs {}",
            simd_result, scalar_result
        );
    }
}

// ============================================================================
// GPU-accelerated batch distance computation — Phase-66
// ============================================================================

/// GPU-accelerated batch distance computation using Candle tensors.
///
/// Computes distances between one query vector and many candidate vectors
/// in a single GPU kernel launch, avoiding per-vector overhead.
///
/// # Architecture
///
/// ```text
///   query [1, D]   ─┐
///                    ├──▶ GPU matmul / batch-op ──▶ distances [N]
///   vectors [N, D] ─┘
/// ```
///
/// Falls back to CPU SIMD when GPU is unavailable or data is too small.
#[cfg(feature = "candle")]
pub mod gpu_batch {
    use super::*;
    use candle_core::{Device, Tensor, DType};
    use crate::device::{
        DeviceSelector, WorkloadHint, GpuFallbackEvent,
        FallbackReason, DevicePreference, global_selector,
    };

    /// Result of a GPU batch distance computation.
    pub struct BatchDistanceResult {
        /// Distances from query to each candidate (length = N).
        pub distances: Vec<f32>,
        /// Whether GPU was used for this computation.
        pub used_gpu: bool,
    }

    /// Compute distances between a query vector and N candidate vectors.
    ///
    /// Automatically selects GPU or CPU based on `DeviceSelector` and data size.
    ///
    /// # Arguments
    /// * `query` — query vector [D]
    /// * `vectors` — candidate vectors, packed row-major [N * D]
    /// * `n_vectors` — number of candidate vectors
    /// * `dim` — vector dimension
    /// * `metric` — distance metric to use
    ///
    /// # Returns
    /// `BatchDistanceResult` with distances and GPU usage flag.
    pub fn batch_distances(
        query: &[f32],
        vectors: &[f32],
        n_vectors: usize,
        dim: usize,
        metric: DistanceMetric,
    ) -> BatchDistanceResult {
        batch_distances_with_selector(query, vectors, n_vectors, dim, metric, &global_selector())
    }

    /// Compute batch distances with an explicit `DeviceSelector`.
    pub fn batch_distances_with_selector(
        query: &[f32],
        vectors: &[f32],
        n_vectors: usize,
        dim: usize,
        metric: DistanceMetric,
        selector: &DeviceSelector,
    ) -> BatchDistanceResult {
        assert_eq!(query.len(), dim, "query dimension mismatch");
        assert_eq!(vectors.len(), n_vectors * dim, "vectors buffer size mismatch");

        if n_vectors == 0 {
            return BatchDistanceResult {
                distances: Vec::new(),
                used_gpu: false,
            };
        }

        // Check if GPU is worthwhile
        let hint = WorkloadHint::with_operation(n_vectors, dim, "brute_force_search");

        if selector.should_use_gpu(&hint) {
            match gpu_batch_inner(query, vectors, n_vectors, dim, metric, selector) {
                Ok(distances) => {
                    return BatchDistanceResult {
                        distances,
                        used_gpu: true,
                    };
                }
                Err(e) => {
                    tracing::warn!(
                        error = %e,
                        n_vectors = n_vectors,
                        dim = dim,
                        "GPU batch distance failed, falling back to CPU SIMD"
                    );
                    selector.emit_fallback_public(GpuFallbackEvent {
                        reason: FallbackReason::RuntimeError,
                        detail: format!("GPU batch distance failed: {}", e),
                        operation: "brute_force_search".to_string(),
                        requested_device: DevicePreference::Auto,
                    });
                }
            }
        }

        // CPU SIMD fallback
        let calc = DistanceCalculator::new(metric);
        let distances: Vec<f32> = (0..n_vectors)
            .map(|i| {
                let v = &vectors[i * dim..(i + 1) * dim];
                calc.calculate(query, v)
            })
            .collect();

        BatchDistanceResult {
            distances,
            used_gpu: false,
        }
    }

    /// Inner GPU computation using Candle tensors.
    fn gpu_batch_inner(
        query: &[f32],
        vectors: &[f32],
        n_vectors: usize,
        dim: usize,
        metric: DistanceMetric,
        selector: &DeviceSelector,
    ) -> Result<Vec<f32>, candle_core::Error> {
        let (device, _is_gpu) = selector.candle_device_with_fallback("brute_force_search");

        match metric {
            DistanceMetric::Cosine => {
                gpu_cosine_distances(query, vectors, n_vectors, dim, &device)
            }
            DistanceMetric::Euclidean => {
                gpu_euclidean_distances(query, vectors, n_vectors, dim, &device)
            }
            DistanceMetric::DotProduct => {
                gpu_dot_product_distances(query, vectors, n_vectors, dim, &device)
            }
        }
    }

    /// GPU cosine distance: 1 - (q · v) / (|q| * |v|)
    ///
    /// Optimized: normalize all vectors first, then distance = 1 - matmul.
    fn gpu_cosine_distances(
        query: &[f32],
        vectors: &[f32],
        n_vectors: usize,
        dim: usize,
        device: &Device,
    ) -> Result<Vec<f32>, candle_core::Error> {
        // Create tensors on device
        let q = Tensor::from_vec(query.to_vec(), (1, dim), device)?; // [1, D]
        let v = Tensor::from_vec(vectors.to_vec(), (n_vectors, dim), device)?; // [N, D]

        // L2 normalize
        let q_norm = q.broadcast_div(&q.sqr()?.sum_keepdim(1)?.sqrt()?.clamp(1e-12, f32::MAX)?)?;
        let v_norms = v.sqr()?.sum_keepdim(1)?.sqrt()?.clamp(1e-12, f32::MAX)?;
        let v_norm = v.broadcast_div(&v_norms)?;

        // Cosine similarity via matmul: [1, D] * [D, N] = [1, N]
        let similarities = q_norm.matmul(&v_norm.t()?)?; // [1, N]

        // Distance = 1 - similarity
        let ones = Tensor::ones(similarities.shape(), DType::F32, device)?;
        let distances = (ones - similarities)?;

        // Transfer back to CPU
        let distances = distances
            .to_device(&Device::Cpu)?
            .squeeze(0)?
            .to_vec1::<f32>()?;

        Ok(distances)
    }

    /// GPU Euclidean distance: sqrt(sum((q - v)^2))
    ///
    /// Optimized: |q-v|^2 = |q|^2 + |v|^2 - 2*q·v
    fn gpu_euclidean_distances(
        query: &[f32],
        vectors: &[f32],
        n_vectors: usize,
        dim: usize,
        device: &Device,
    ) -> Result<Vec<f32>, candle_core::Error> {
        let q = Tensor::from_vec(query.to_vec(), (1, dim), device)?;
        let v = Tensor::from_vec(vectors.to_vec(), (n_vectors, dim), device)?;

        // |q|^2 = scalar
        let q_sq = q.sqr()?.sum_keepdim(1)?; // [1, 1]
        // |v_i|^2 for each vector
        let v_sq = v.sqr()?.sum_keepdim(1)?; // [N, 1]
        // q · v_i
        let qv = q.matmul(&v.t()?)?; // [1, N]

        // |q-v|^2 = |q|^2 + |v|^2 - 2*q·v
        let v_sq_t = v_sq.t()?; // [1, N]
        let sq_dists = (q_sq.broadcast_add(&v_sq_t)? - (qv * 2.0)?)?;

        // Clamp to avoid negative due to floating-point errors, then sqrt
        let sq_dists = sq_dists.clamp(0.0, f32::MAX)?;
        let distances = sq_dists.sqrt()?;

        let distances = distances
            .to_device(&Device::Cpu)?
            .squeeze(0)?
            .to_vec1::<f32>()?;

        Ok(distances)
    }

    /// GPU dot product distance: -dot(q, v)
    ///
    /// Negative because smaller distance = more similar.
    fn gpu_dot_product_distances(
        query: &[f32],
        vectors: &[f32],
        n_vectors: usize,
        dim: usize,
        device: &Device,
    ) -> Result<Vec<f32>, candle_core::Error> {
        let q = Tensor::from_vec(query.to_vec(), (1, dim), device)?;
        let v = Tensor::from_vec(vectors.to_vec(), (n_vectors, dim), device)?;

        let dots = q.matmul(&v.t()?)?; // [1, N]
        let neg_dots = dots.neg()?;

        let distances = neg_dots
            .to_device(&Device::Cpu)?
            .squeeze(0)?
            .to_vec1::<f32>()?;

        Ok(distances)
    }

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

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

        fn approx_eq(a: f32, b: f32, eps: f32) -> bool {
            (a - b).abs() < eps
        }

        #[test]
        fn test_gpu_batch_empty() {
            let query = vec![1.0, 0.0, 0.0];
            let result = batch_distances(&query, &[], 0, 3, DistanceMetric::Cosine);
            assert!(result.distances.is_empty());
            assert!(!result.used_gpu);
        }

        #[test]
        fn test_gpu_batch_cosine_vs_cpu() {
            let dim = 128;
            let n = 100;
            let query: Vec<f32> = (0..dim).map(|i| (i as f32 * 0.1).sin()).collect();
            let mut vectors = Vec::with_capacity(n * dim);
            for i in 0..n {
                for d in 0..dim {
                    vectors.push(((i * dim + d) as f32 * 0.07).cos());
                }
            }

            // CPU reference
            let calc = DistanceCalculator::cosine();
            let cpu_dists: Vec<f32> = (0..n)
                .map(|i| calc.calculate(&query, &vectors[i * dim..(i + 1) * dim]))
                .collect();

            // GPU batch (may fallback to CPU if no GPU)
            let gpu_result = batch_distances(&query, &vectors, n, dim, DistanceMetric::Cosine);

            assert_eq!(gpu_result.distances.len(), n);
            for (i, (g, c)) in gpu_result.distances.iter().zip(cpu_dists.iter()).enumerate() {
                assert!(
                    approx_eq(*g, *c, 1e-3),
                    "Cosine mismatch at index {}: gpu={} cpu={}",
                    i, g, c
                );
            }
        }

        #[test]
        fn test_gpu_batch_euclidean_vs_cpu() {
            let dim = 64;
            let n = 50;
            let query: Vec<f32> = (0..dim).map(|i| i as f32 * 0.1).collect();
            let mut vectors = Vec::with_capacity(n * dim);
            for i in 0..n {
                for d in 0..dim {
                    vectors.push((i as f32 + d as f32) * 0.05);
                }
            }

            let calc = DistanceCalculator::euclidean();
            let cpu_dists: Vec<f32> = (0..n)
                .map(|i| calc.calculate(&query, &vectors[i * dim..(i + 1) * dim]))
                .collect();

            let gpu_result = batch_distances(&query, &vectors, n, dim, DistanceMetric::Euclidean);

            assert_eq!(gpu_result.distances.len(), n);
            for (i, (g, c)) in gpu_result.distances.iter().zip(cpu_dists.iter()).enumerate() {
                assert!(
                    approx_eq(*g, *c, 1e-2),
                    "Euclidean mismatch at index {}: gpu={} cpu={}",
                    i, g, c
                );
            }
        }

        #[test]
        fn test_gpu_batch_dot_product_vs_cpu() {
            let dim = 32;
            let n = 20;
            let query: Vec<f32> = (0..dim).map(|i| (i as f32).sin()).collect();
            let mut vectors = Vec::with_capacity(n * dim);
            for i in 0..n {
                for d in 0..dim {
                    vectors.push(((i + d) as f32 * 0.3).cos());
                }
            }

            let calc = DistanceCalculator::dot_product();
            let cpu_dists: Vec<f32> = (0..n)
                .map(|i| calc.calculate(&query, &vectors[i * dim..(i + 1) * dim]))
                .collect();

            let gpu_result = batch_distances(&query, &vectors, n, dim, DistanceMetric::DotProduct);

            assert_eq!(gpu_result.distances.len(), n);
            for (i, (g, c)) in gpu_result.distances.iter().zip(cpu_dists.iter()).enumerate() {
                assert!(
                    approx_eq(*g, *c, 1e-3),
                    "DotProduct mismatch at index {}: gpu={} cpu={}",
                    i, g, c
                );
            }
        }

        #[test]
        fn test_gpu_batch_with_cpu_selector() {
            use crate::device::{GpuAccelConfig, DevicePreference};
            let selector = DeviceSelector::new(GpuAccelConfig {
                device_preference: DevicePreference::Cpu,
                ..Default::default()
            });

            let dim = 16;
            let n = 10;
            let query: Vec<f32> = (0..dim).map(|i| i as f32).collect();
            let vectors: Vec<f32> = (0..n * dim).map(|i| i as f32 * 0.1).collect();

            let result = batch_distances_with_selector(
                &query, &vectors, n, dim, DistanceMetric::Cosine, &selector,
            );
            assert_eq!(result.distances.len(), n);
            assert!(!result.used_gpu, "Should use CPU when preference is Cpu");
        }

        #[test]
        fn test_gpu_batch_large_scale_consistency() {
            // Larger test to verify stability
            let dim = 384;
            let n = 500;
            let query: Vec<f32> = (0..dim).map(|i| (i as f32 * 0.01).sin()).collect();
            let mut vectors = Vec::with_capacity(n * dim);
            for i in 0..n {
                for d in 0..dim {
                    vectors.push(((i * 13 + d * 7) as f32 * 0.001).cos());
                }
            }

            let calc = DistanceCalculator::cosine();
            let cpu_dists: Vec<f32> = (0..n)
                .map(|i| calc.calculate(&query, &vectors[i * dim..(i + 1) * dim]))
                .collect();

            let gpu_result = batch_distances(&query, &vectors, n, dim, DistanceMetric::Cosine);
            assert_eq!(gpu_result.distances.len(), n);

            // Verify top-k ordering is consistent
            let mut cpu_indexed: Vec<(usize, f32)> =
                cpu_dists.iter().copied().enumerate().collect();
            let mut gpu_indexed: Vec<(usize, f32)> =
                gpu_result.distances.iter().copied().enumerate().collect();
            cpu_indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
            gpu_indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

            // Top-10 should match
            let cpu_top10: Vec<usize> = cpu_indexed.iter().take(10).map(|(i, _)| *i).collect();
            let gpu_top10: Vec<usize> = gpu_indexed.iter().take(10).map(|(i, _)| *i).collect();
            assert_eq!(cpu_top10, gpu_top10, "Top-10 ordering should match");
        }
    }
}