//! Vector Quantization for compression
//!
//! This module provides two quantization methods:
//! 
//! ## Scalar Quantization (SQ)
//! - Compresses f32 to u8, achieving 4x compression
//! - Higher accuracy, lower compression ratio
//! - Best for precision-sensitive applications
//! 
//! ## Product Quantization (PQ)
//! - Compresses vectors to ~1/32 of original size
//! - Lower accuracy, higher compression ratio
//! - Best for memory-constrained applications

use rand::seq::SliceRandom;
use rand::Rng;
use serde::{Deserialize, Serialize};

use crate::device::{DeviceSelector, WorkloadHint};

// ============================================================================
// Scalar Quantization
// ============================================================================

/// Configuration for Scalar Quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SQConfig {
    /// Whether to use per-dimension min/max (true) or global min/max (false)
    pub per_dimension: bool,
}

impl Default for SQConfig {
    fn default() -> Self {
        Self {
            per_dimension: true,
        }
    }
}

/// Scalar Quantizer for vector compression
/// 
/// Compresses f32 vectors to u8 vectors, achieving 4x compression ratio.
/// Each dimension is independently scaled to [0, 255] range.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScalarQuantizer {
    /// Configuration
    config: SQConfig,
    
    /// Vector dimension
    dimension: usize,
    
    /// Minimum values per dimension (or single global min)
    min_vals: Vec<f32>,
    
    /// Maximum values per dimension (or single global max)
    max_vals: Vec<f32>,
    
    /// Scale factors for quantization: 255 / (max - min)
    scales: Vec<f32>,
    
    /// Whether the quantizer has been trained
    trained: bool,
}

impl ScalarQuantizer {
    /// Create a new untrained Scalar Quantizer
    pub fn new(dimension: usize, config: SQConfig) -> Self {
        Self {
            config,
            dimension,
            min_vals: Vec::new(),
            max_vals: Vec::new(),
            scales: Vec::new(),
            trained: false,
        }
    }
    
    /// Create with default config
    pub fn for_dimension(dimension: usize) -> Self {
        Self::new(dimension, SQConfig::default())
    }
    
    /// Check if trained
    pub fn is_trained(&self) -> bool {
        self.trained
    }
    
    /// Get compression ratio (always 4x for f32 -> u8)
    pub fn compression_ratio(&self) -> f32 {
        4.0
    }
    
    /// Get code size in bytes
    pub fn code_size(&self) -> usize {
        self.dimension
    }
    
    /// Train the quantizer on a set of vectors
    pub fn train(&mut self, vectors: &[Vec<f32>]) -> Result<(), SQError> {
        self.train_with_selector(vectors, None)
    }

    /// Train the quantizer with explicit DeviceSelector for GPU acceleration.
    ///
    /// When `selector` is provided and GPU is available, uses GPU-accelerated
    /// min/max reduction for large datasets. Falls back to CPU seamlessly.
    pub fn train_with_selector(
        &mut self,
        vectors: &[Vec<f32>],
        selector: Option<&DeviceSelector>,
    ) -> Result<(), SQError> {
        if vectors.is_empty() {
            return Err(SQError::NotEnoughData);
        }
        
        // Validate dimensions
        for v in vectors {
            if v.len() != self.dimension {
                return Err(SQError::DimensionMismatch {
                    expected: self.dimension,
                    got: v.len(),
                });
            }
        }

        // Check if GPU acceleration is worthwhile
        let sel = selector.cloned().unwrap_or_else(DeviceSelector::auto);
        let hint = WorkloadHint::with_operation(vectors.len(), self.dimension, "sq_train");
        let use_gpu = sel.should_use_gpu(&hint);

        // Try GPU path for per-dimension min/max computation
        #[cfg(feature = "candle")]
        if use_gpu && self.config.per_dimension {
            if let Ok(()) = self.train_gpu_minmax(vectors, &sel) {
                return Ok(());
            }
            // GPU failed, fall through to CPU
        }
        
        if self.config.per_dimension {
            // Per-dimension min/max (CPU)
            self.min_vals = vec![f32::MAX; self.dimension];
            self.max_vals = vec![f32::MIN; self.dimension];
            
            for v in vectors {
                for (i, &val) in v.iter().enumerate() {
                    if val < self.min_vals[i] {
                        self.min_vals[i] = val;
                    }
                    if val > self.max_vals[i] {
                        self.max_vals[i] = val;
                    }
                }
            }
        } else {
            // Global min/max
            let mut global_min = f32::MAX;
            let mut global_max = f32::MIN;
            
            for v in vectors {
                for &val in v {
                    if val < global_min {
                        global_min = val;
                    }
                    if val > global_max {
                        global_max = val;
                    }
                }
            }
            
            self.min_vals = vec![global_min; self.dimension];
            self.max_vals = vec![global_max; self.dimension];
        }
        
        // Compute scale factors
        self.scales = Vec::with_capacity(self.dimension);
        for i in 0..self.dimension {
            let range = self.max_vals[i] - self.min_vals[i];
            if range > 1e-10 {
                self.scales.push(255.0 / range);
            } else {
                // Handle constant dimension
                self.scales.push(1.0);
            }
        }
        
        self.trained = true;
        Ok(())
    }

    /// GPU-accelerated min/max computation for SQ training.
    ///
    /// Uses candle tensor operations for parallel reduction across all dimensions.
    #[cfg(feature = "candle")]
    fn train_gpu_minmax(
        &mut self,
        vectors: &[Vec<f32>],
        selector: &DeviceSelector,
    ) -> Result<(), SQError> {
        use candle_core::Tensor;

        let n_vectors = vectors.len();
        let dimension = self.dimension;

        let (device, _is_gpu) = selector.candle_device_with_fallback("sq_train_minmax");

        // Flatten vectors into contiguous buffer
        let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();

        // Transfer to device and compute min/max
        let tensor = Tensor::from_slice(&flat, (n_vectors, dimension), &device)
            .map_err(|e| {
                tracing::warn!("SQ GPU: tensor creation failed: {e}");
                SQError::NotEnoughData // fallback trigger
            })?;

        let min_tensor = tensor.min(0).map_err(|e| {
            tracing::warn!("SQ GPU: min reduction failed: {e}");
            SQError::NotEnoughData
        })?;
        let max_tensor = tensor.max(0).map_err(|e| {
            tracing::warn!("SQ GPU: max reduction failed: {e}");
            SQError::NotEnoughData
        })?;

        // Transfer back to CPU
        self.min_vals = min_tensor
            .to_device(&candle_core::Device::Cpu)
            .and_then(|t| t.to_vec1::<f32>())
            .map_err(|e| {
                tracing::warn!("SQ GPU: min copy failed: {e}");
                SQError::NotEnoughData
            })?;

        self.max_vals = max_tensor
            .to_device(&candle_core::Device::Cpu)
            .and_then(|t| t.to_vec1::<f32>())
            .map_err(|e| {
                tracing::warn!("SQ GPU: max copy failed: {e}");
                SQError::NotEnoughData
            })?;

        // Compute scale factors
        self.scales = Vec::with_capacity(dimension);
        for i in 0..dimension {
            let range = self.max_vals[i] - self.min_vals[i];
            if range > 1e-10 {
                self.scales.push(255.0 / range);
            } else {
                self.scales.push(1.0);
            }
        }

        self.trained = true;
        tracing::info!(
            "SQ GPU training completed: {}d, {} vectors",
            dimension, n_vectors
        );
        Ok(())
    }
    
    /// Quantize a vector to u8 codes
    pub fn quantize(&self, vector: &[f32]) -> Result<Vec<u8>, SQError> {
        if !self.trained {
            return Err(SQError::NotTrained);
        }
        
        if vector.len() != self.dimension {
            return Err(SQError::DimensionMismatch {
                expected: self.dimension,
                got: vector.len(),
            });
        }
        
        let codes: Vec<u8> = vector.iter()
            .enumerate()
            .map(|(i, &v)| {
                let normalized = (v - self.min_vals[i]) * self.scales[i];
                normalized.clamp(0.0, 255.0) as u8
            })
            .collect();
        
        Ok(codes)
    }
    
    /// Dequantize u8 codes back to approximate f32 vector
    pub fn dequantize(&self, codes: &[u8]) -> Result<Vec<f32>, SQError> {
        if !self.trained {
            return Err(SQError::NotTrained);
        }
        
        if codes.len() != self.dimension {
            return Err(SQError::DimensionMismatch {
                expected: self.dimension,
                got: codes.len(),
            });
        }
        
        let vector: Vec<f32> = codes.iter()
            .enumerate()
            .map(|(i, &code)| {
                self.min_vals[i] + (code as f32) / self.scales[i]
            })
            .collect();
        
        Ok(vector)
    }
    
    /// Compute squared Euclidean distance between query and quantized vector
    pub fn distance_sq(&self, query: &[f32], codes: &[u8]) -> Result<f32, SQError> {
        if !self.trained {
            return Err(SQError::NotTrained);
        }
        
        if query.len() != self.dimension {
            return Err(SQError::DimensionMismatch {
                expected: self.dimension,
                got: query.len(),
            });
        }
        
        let mut dist_sq = 0.0f32;
        for (i, (&q, &code)) in query.iter().zip(codes.iter()).enumerate() {
            let decoded = self.min_vals[i] + (code as f32) / self.scales[i];
            let diff = q - decoded;
            dist_sq += diff * diff;
        }
        
        Ok(dist_sq)
    }
    
    /// Compute Euclidean distance
    pub fn distance(&self, query: &[f32], codes: &[u8]) -> Result<f32, SQError> {
        Ok(self.distance_sq(query, codes)?.sqrt())
    }
    
    /// Batch quantize multiple vectors
    pub fn quantize_batch(&self, vectors: &[Vec<f32>]) -> Result<Vec<Vec<u8>>, SQError> {
        vectors.iter().map(|v| self.quantize(v)).collect()
    }
    
    /// Get dimension
    pub fn dimension(&self) -> usize {
        self.dimension
    }
}

/// Scalar Quantization errors
#[derive(Debug, thiserror::Error)]
pub enum SQError {
    #[error("Quantizer not trained")]
    NotTrained,
    
    #[error("Not enough training data")]
    NotEnoughData,
    
    #[error("Dimension mismatch: expected {expected}, got {got}")]
    DimensionMismatch { expected: usize, got: usize },
}

// ============================================================================
// Product Quantization
// ============================================================================

/// Configuration for Product Quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PQConfig {
    /// Number of subspaces (subvectors)
    /// Total dimension must be divisible by this
    pub num_subspaces: usize,
    
    /// Number of centroids per subspace (typically 256 for 8-bit codes)
    pub num_centroids: usize,
    
    /// Number of k-means iterations for training
    pub training_iterations: usize,
    
    /// Sample size for training (0 = use all vectors)
    pub training_sample_size: usize,
}

impl Default for PQConfig {
    fn default() -> Self {
        Self {
            num_subspaces: 192,      // For 1536-dim vectors: 1536/192 = 8 dims per subspace
            num_centroids: 256,       // 8-bit codes
            training_iterations: 25,
            training_sample_size: 10000,
        }
    }
}

impl PQConfig {
    /// Create PQ config for a specific dimension
    pub fn for_dimension(dimension: usize) -> Self {
        // Choose num_subspaces such that sub_dimension is reasonable (4-16)
        let num_subspaces = if dimension % 8 == 0 {
            dimension / 8
        } else if dimension % 4 == 0 {
            dimension / 4
        } else {
            // Find a divisor
            for d in [8, 4, 2, 1].iter() {
                if dimension % d == 0 {
                    return Self {
                        num_subspaces: dimension / d,
                        ..Default::default()
                    };
                }
            }
            dimension // Worst case: 1-dim subspaces
        };
        
        Self {
            num_subspaces,
            ..Default::default()
        }
    }
    
    /// Calculate sub-dimension
    pub fn sub_dimension(&self, total_dimension: usize) -> usize {
        total_dimension / self.num_subspaces
    }
}

/// Product Quantizer for vector compression
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProductQuantizer {
    /// Configuration
    config: PQConfig,
    
    /// Vector dimension
    dimension: usize,
    
    /// Dimension of each subspace
    sub_dimension: usize,
    
    /// Centroids for each subspace: [subspace][centroid][dim]
    centroids: Vec<Vec<Vec<f32>>>,
    
    /// Whether the quantizer has been trained
    trained: bool,
}

impl ProductQuantizer {
    /// Create a new untrained Product Quantizer
    pub fn new(dimension: usize, config: PQConfig) -> Self {
        assert!(
            dimension % config.num_subspaces == 0,
            "Dimension {} must be divisible by num_subspaces {}",
            dimension, config.num_subspaces
        );
        
        let sub_dimension = dimension / config.num_subspaces;
        
        Self {
            config,
            dimension,
            sub_dimension,
            centroids: Vec::new(),
            trained: false,
        }
    }
    
    /// Create with default config for dimension
    pub fn for_dimension(dimension: usize) -> Self {
        Self::new(dimension, PQConfig::for_dimension(dimension))
    }
    
    /// Check if the quantizer is trained
    pub fn is_trained(&self) -> bool {
        self.trained
    }
    
    /// Get the code size in bytes
    pub fn code_size(&self) -> usize {
        self.config.num_subspaces
    }
    
    /// Get compression ratio compared to f32 vectors
    pub fn compression_ratio(&self) -> f32 {
        (self.dimension * std::mem::size_of::<f32>()) as f32 / self.code_size() as f32
    }
    
    /// Train the quantizer on a set of vectors
    pub fn train(&mut self, vectors: &[Vec<f32>]) -> Result<(), PQError> {
        self.train_with_selector(vectors, None)
    }

    /// Train the quantizer with explicit DeviceSelector for GPU acceleration.
    ///
    /// When GPU is available and dataset is large enough, uses GPU KMeans
    /// for each subspace codebook training. Falls back to CPU seamlessly.
    pub fn train_with_selector(
        &mut self,
        vectors: &[Vec<f32>],
        selector: Option<&DeviceSelector>,
    ) -> Result<(), PQError> {
        if vectors.is_empty() {
            return Err(PQError::NotEnoughData);
        }
        
        // Validate dimensions
        for v in vectors {
            if v.len() != self.dimension {
                return Err(PQError::DimensionMismatch {
                    expected: self.dimension,
                    got: v.len(),
                });
            }
        }
        
        // Sample vectors if needed
        let training_vectors: Vec<&Vec<f32>> = if self.config.training_sample_size > 0 
            && vectors.len() > self.config.training_sample_size 
        {
            let mut rng = rand::thread_rng();
            vectors.choose_multiple(&mut rng, self.config.training_sample_size).collect()
        } else {
            vectors.iter().collect()
        };

        // Check if GPU acceleration is worthwhile
        let sel = selector.cloned().unwrap_or_else(DeviceSelector::auto);
        let hint = WorkloadHint::with_operation(
            training_vectors.len(), self.sub_dimension, "pq_kmeans_train",
        );
        let use_gpu = sel.should_use_gpu(&hint);
        
        // Initialize centroids for each subspace
        self.centroids = Vec::with_capacity(self.config.num_subspaces);
        
        for subspace_idx in 0..self.config.num_subspaces {
            let start = subspace_idx * self.sub_dimension;
            let end = start + self.sub_dimension;
            
            // Extract subvectors for this subspace
            let subvectors: Vec<Vec<f32>> = training_vectors
                .iter()
                .map(|v| v[start..end].to_vec())
                .collect();

            // Try GPU KMeans path when available
            #[cfg(feature = "gpu-index")]
            if use_gpu {
                if let Some(gpu_centroids) = self.kmeans_gpu(
                    &subvectors, subspace_idx, &sel,
                ) {
                    self.centroids.push(gpu_centroids);
                    continue;
                }
                // GPU failed, fall through to CPU
            }
            
            // CPU k-means to find centroids
            let centroids = self.kmeans(
                &subvectors,
                self.config.num_centroids,
                self.config.training_iterations,
            );
            
            self.centroids.push(centroids);
        }
        
        self.trained = true;
        Ok(())
    }

    /// GPU-accelerated KMeans for a single PQ subspace.
    ///
    /// Returns `Some(centroids)` on success, `None` on failure (caller falls back to CPU).
    #[cfg(feature = "gpu-index")]
    fn kmeans_gpu(
        &self,
        subvectors: &[Vec<f32>],
        subspace_idx: usize,
        selector: &DeviceSelector,
    ) -> Option<Vec<Vec<f32>>> {
        use crate::gpu::kmeans::{GpuKMeansConfig, GpuKMeansTrainer};

        let n_vectors = subvectors.len();
        let sub_dim = self.sub_dimension;

        // Flatten subvectors into contiguous buffer for GPU trainer
        let flat: Vec<f32> = subvectors.iter().flat_map(|v| v.iter().copied()).collect();

        let kmeans_config = GpuKMeansConfig {
            n_clusters: self.config.num_centroids,
            max_iterations: self.config.training_iterations,
            tolerance: 1e-4,
            seed: 42 + subspace_idx as u64,
            use_kmeans_plus_plus: true,
            verbose: false,
            ..Default::default()
        };

        let trainer = GpuKMeansTrainer::new(kmeans_config);
        match trainer.train_with_selector(&flat, n_vectors, sub_dim, None, Some(selector)) {
            Ok(result) => {
                if subspace_idx == 0 {
                    tracing::info!(
                        "PQ subspace GPU KMeans: {} vectors × {}d, {} iters, gpu={}",
                        n_vectors, sub_dim, result.iterations, result.used_gpu,
                    );
                }
                // Convert flat centroids [n_clusters × sub_dim] to Vec<Vec<f32>>
                let centroids: Vec<Vec<f32>> = (0..self.config.num_centroids)
                    .map(|c| {
                        let s = c * sub_dim;
                        result.centroids[s..s + sub_dim].to_vec()
                    })
                    .collect();
                Some(centroids)
            }
            Err(e) => {
                tracing::warn!(
                    "PQ subspace {} GPU KMeans failed, falling back to CPU: {}",
                    subspace_idx, e,
                );
                None
            }
        }
    }
    
    /// K-means clustering for a single subspace
    fn kmeans(&self, vectors: &[Vec<f32>], k: usize, iterations: usize) -> Vec<Vec<f32>> {
        let dim = self.sub_dimension;
        let n = vectors.len();
        
        if n == 0 || k == 0 {
            return vec![vec![0.0; dim]; k];
        }
        
        let mut rng = rand::thread_rng();
        
        // Initialize centroids - sample k vectors or pad with zeros
        let mut centroids: Vec<Vec<f32>> = Vec::with_capacity(k);
        
        // Use simple random sampling for initialization
        let sample_size = k.min(n);
        let mut used_indices = std::collections::HashSet::new();
        
        while centroids.len() < sample_size {
            let idx = rng.gen_range(0..n);
            if !used_indices.contains(&idx) {
                used_indices.insert(idx);
                centroids.push(vectors[idx].clone());
            }
            // Prevent infinite loop if not enough unique vectors
            if used_indices.len() >= n {
                break;
            }
        }
        
        // Pad with random vectors from the dataset if needed
        while centroids.len() < k {
            let idx = rng.gen_range(0..n);
            centroids.push(vectors[idx].clone());
        }
        
        // Ensure exactly k centroids
        centroids.truncate(k);
        
        // K-means iterations
        let mut assignments = vec![0usize; n];
        
        for _ in 0..iterations {
            // Assign each vector to nearest centroid
            for (i, v) in vectors.iter().enumerate() {
                let mut min_dist = f32::MAX;
                let mut min_idx = 0usize;
                
                for (j, c) in centroids.iter().enumerate() {
                    let dist = squared_euclidean(v, c);
                    if dist < min_dist {
                        min_dist = dist;
                        min_idx = j;
                    }
                }
                
                // Ensure min_idx is within bounds
                assignments[i] = min_idx.min(k - 1);
            }
            
            // Update centroids
            let mut new_centroids = vec![vec![0.0; dim]; k];
            let mut counts = vec![0usize; k];
            
            for (i, v) in vectors.iter().enumerate() {
                let cluster = assignments[i];
                if cluster < k {
                    counts[cluster] += 1;
                    for (j, &val) in v.iter().enumerate() {
                        new_centroids[cluster][j] += val;
                    }
                }
            }
            
            for (i, c) in new_centroids.iter_mut().enumerate() {
                if counts[i] > 0 {
                    for val in c.iter_mut() {
                        *val /= counts[i] as f32;
                    }
                } else {
                    // Keep old centroid for empty clusters
                    *c = centroids[i].clone();
                }
            }
            
            centroids = new_centroids;
        }
        
        centroids
    }
    
    /// Encode a vector to PQ codes
    pub fn encode(&self, vector: &[f32]) -> Result<Vec<u8>, PQError> {
        if !self.trained {
            return Err(PQError::NotTrained);
        }
        
        if vector.len() != self.dimension {
            return Err(PQError::DimensionMismatch {
                expected: self.dimension,
                got: vector.len(),
            });
        }
        
        let mut codes = Vec::with_capacity(self.config.num_subspaces);
        
        for (subspace_idx, subspace_centroids) in self.centroids.iter().enumerate() {
            let start = subspace_idx * self.sub_dimension;
            let end = start + self.sub_dimension;
            let subvector = &vector[start..end];
            
            // Find nearest centroid
            let mut min_dist = f32::MAX;
            let mut min_idx = 0u8;
            
            for (i, centroid) in subspace_centroids.iter().enumerate() {
                let dist = squared_euclidean(subvector, centroid);
                if dist < min_dist {
                    min_dist = dist;
                    min_idx = i as u8;
                }
            }
            
            codes.push(min_idx);
        }
        
        Ok(codes)
    }
    
    /// Decode PQ codes back to an approximate vector
    pub fn decode(&self, codes: &[u8]) -> Result<Vec<f32>, PQError> {
        if !self.trained {
            return Err(PQError::NotTrained);
        }
        
        if codes.len() != self.config.num_subspaces {
            return Err(PQError::InvalidCodeLength {
                expected: self.config.num_subspaces,
                got: codes.len(),
            });
        }
        
        let mut vector = Vec::with_capacity(self.dimension);
        
        for (subspace_idx, &code) in codes.iter().enumerate() {
            let centroid = &self.centroids[subspace_idx][code as usize];
            vector.extend_from_slice(centroid);
        }
        
        Ok(vector)
    }
    
    /// Compute asymmetric distance between a query vector and encoded vector
    /// 
    /// This is more accurate than symmetric distance (comparing two encoded vectors)
    pub fn asymmetric_distance(&self, query: &[f32], codes: &[u8]) -> Result<f32, PQError> {
        if !self.trained {
            return Err(PQError::NotTrained);
        }
        
        if query.len() != self.dimension {
            return Err(PQError::DimensionMismatch {
                expected: self.dimension,
                got: query.len(),
            });
        }
        
        if codes.len() != self.config.num_subspaces {
            return Err(PQError::InvalidCodeLength {
                expected: self.config.num_subspaces,
                got: codes.len(),
            });
        }
        
        let mut total_dist = 0.0f32;
        
        for (subspace_idx, &code) in codes.iter().enumerate() {
            let start = subspace_idx * self.sub_dimension;
            let end = start + self.sub_dimension;
            let subquery = &query[start..end];
            let centroid = &self.centroids[subspace_idx][code as usize];
            
            total_dist += squared_euclidean(subquery, centroid);
        }
        
        Ok(total_dist.sqrt())
    }
    
    /// Precompute distance table for a query
    /// 
    /// This allows very fast distance computation when searching multiple codes.
    ///
    /// Phase-67: When the `candle` feature is enabled, uses GPU-accelerated
    /// batch matmul to compute all subspace distances in parallel.
    /// Falls back to CPU if GPU is unavailable or errors occur.
    pub fn compute_distance_table(&self, query: &[f32]) -> Result<Vec<Vec<f32>>, PQError> {
        if !self.trained {
            return Err(PQError::NotTrained);
        }
        
        if query.len() != self.dimension {
            return Err(PQError::DimensionMismatch {
                expected: self.dimension,
                got: query.len(),
            });
        }

        #[cfg(feature = "candle")]
        {
            let hint = WorkloadHint::with_operation(
                self.config.num_subspaces * self.config.num_centroids,
                self.sub_dimension,
                "pq_distance_table",
            );
            let selector = crate::device::global_selector();
            if selector.should_use_gpu(&hint) {
                match self.compute_distance_table_gpu(&selector, query) {
                    Ok(table) => return Ok(table),
                    Err(e) => {
                        tracing::warn!(
                            error = %e,
                            num_subspaces = self.config.num_subspaces,
                            num_centroids = self.config.num_centroids,
                            sub_dim = self.sub_dimension,
                            "GPU PQ distance table failed, falling back to CPU"
                        );
                        selector.emit_fallback_public(crate::device::GpuFallbackEvent {
                            reason: crate::device::FallbackReason::RuntimeError,
                            detail: format!("GPU PQ distance table: {}", e),
                            operation: "pq_distance_table".to_string(),
                            requested_device: crate::device::DevicePreference::Auto,
                        });
                    }
                }
            }
        }
        
        self.compute_distance_table_cpu(query)
    }

    /// CPU implementation of distance table computation.
    fn compute_distance_table_cpu(&self, query: &[f32]) -> Result<Vec<Vec<f32>>, PQError> {
        let mut table = Vec::with_capacity(self.config.num_subspaces);
        
        for (subspace_idx, subspace_centroids) in self.centroids.iter().enumerate() {
            let start = subspace_idx * self.sub_dimension;
            let end = start + self.sub_dimension;
            let subquery = &query[start..end];
            
            let distances: Vec<f32> = subspace_centroids
                .iter()
                .map(|c| squared_euclidean(subquery, c))
                .collect();
            
            table.push(distances);
        }
        
        Ok(table)
    }

    /// Phase-67: GPU-accelerated distance table computation.
    ///
    /// For each subspace, packs the centroid matrix [K, sub_dim] and the
    /// sub-query [1, sub_dim] as Candle tensors, then computes:
    ///   dist² = |q|² + |c|² - 2·q·c
    /// using a single matmul per subspace.
    #[cfg(feature = "candle")]
    fn compute_distance_table_gpu(
        &self,
        selector: &DeviceSelector,
        query: &[f32],
    ) -> Result<Vec<Vec<f32>>, PQError> {
        use candle_core::{Device, Tensor, DType};

        let (device, _is_gpu) = selector.candle_device_with_fallback("pq_distance_table");
        let num_centroids = self.config.num_centroids;
        let sub_dim = self.sub_dimension;

        let mut table = Vec::with_capacity(self.config.num_subspaces);

        for (subspace_idx, subspace_centroids) in self.centroids.iter().enumerate() {
            let start = subspace_idx * sub_dim;
            let end = start + sub_dim;
            let subquery = &query[start..end];

            // Pack centroids into flat buffer [K * sub_dim]
            let k = subspace_centroids.len();
            let mut flat_centroids = Vec::with_capacity(k * sub_dim);
            for c in subspace_centroids {
                flat_centroids.extend_from_slice(c);
            }

            // Build tensors on device
            let q = Tensor::from_vec(subquery.to_vec(), (1, sub_dim), &device)
                .map_err(|e| PQError::GpuComputeError(format!("GPU tensor error: {}", e)))?;
            let c = Tensor::from_vec(flat_centroids, (k, sub_dim), &device)
                .map_err(|e| PQError::GpuComputeError(format!("GPU tensor error: {}", e)))?;

            // |q|² — scalar broadcast to [1, 1]
            let q_sq = q.sqr()
                .and_then(|t| t.sum_keepdim(1))
                .map_err(|e| PQError::GpuComputeError(format!("GPU error: {}", e)))?;

            // |c_i|² for each centroid → [K, 1]
            let c_sq = c.sqr()
                .and_then(|t| t.sum_keepdim(1))
                .map_err(|e| PQError::GpuComputeError(format!("GPU error: {}", e)))?;

            // q · c^T → [1, K]
            let qc = q.matmul(&c.t()
                .map_err(|e| PQError::GpuComputeError(format!("GPU error: {}", e)))?)
                .map_err(|e| PQError::GpuComputeError(format!("GPU error: {}", e)))?;

            // dist² = |q|² + |c|² - 2·q·c
            let c_sq_t = c_sq.t()
                .map_err(|e| PQError::GpuComputeError(format!("GPU error: {}", e)))?;
            let distances = q_sq.broadcast_add(&c_sq_t)
                .and_then(|t| {
                    let two_qc = (qc * 2.0)?;
                    t - two_qc
                })
                .and_then(|t| t.clamp(0.0, f32::MAX))
                .map_err(|e| PQError::GpuComputeError(format!("GPU error: {}", e)))?;

            // Transfer back to CPU
            let dists: Vec<f32> = distances
                .to_device(&Device::Cpu)
                .and_then(|t| t.squeeze(0))
                .and_then(|t| t.to_vec1::<f32>())
                .map_err(|e| PQError::GpuComputeError(format!("GPU transfer error: {}", e)))?;

            table.push(dists);
        }

        Ok(table)
    }
    
    /// Compute distance using precomputed distance table
    /// 
    /// This is very fast: O(num_subspaces) lookups and additions
    pub fn distance_with_table(&self, table: &[Vec<f32>], codes: &[u8]) -> f32 {
        let mut total = 0.0f32;
        
        for (subspace_idx, &code) in codes.iter().enumerate() {
            total += table[subspace_idx][code as usize];
        }
        
        total.sqrt()
    }
    
    /// Batch encode multiple vectors
    pub fn encode_batch(&self, vectors: &[Vec<f32>]) -> Result<Vec<Vec<u8>>, PQError> {
        vectors.iter().map(|v| self.encode(v)).collect()
    }
    
    /// Get the configuration
    pub fn config(&self) -> &PQConfig {
        &self.config
    }
    
    /// Get the dimension
    pub fn dimension(&self) -> usize {
        self.dimension
    }
}

/// Squared Euclidean distance between two slices
#[inline]
fn squared_euclidean(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b.iter())
        .map(|(&x, &y)| {
            let diff = x - y;
            diff * diff
        })
        .sum()
}

/// PQ Errors
#[derive(Debug, thiserror::Error)]
pub enum PQError {
    #[error("Quantizer not trained")]
    NotTrained,
    
    #[error("Not enough training data")]
    NotEnoughData,
    
    #[error("Dimension mismatch: expected {expected}, got {got}")]
    DimensionMismatch { expected: usize, got: usize },
    
    #[error("Invalid code length: expected {expected}, got {got}")]
    InvalidCodeLength { expected: usize, got: usize },

    /// Phase-67: GPU compute error (distance table, etc.)
    #[error("GPU compute error: {0}")]
    GpuComputeError(String),
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::Rng;
    
    fn generate_random_vectors(count: usize, dim: usize) -> Vec<Vec<f32>> {
        let mut rng = rand::thread_rng();
        (0..count)
            .map(|_| (0..dim).map(|_| rng.gen::<f32>()).collect())
            .collect()
    }
    
    // =========================================================================
    // Scalar Quantization Tests
    // =========================================================================
    
    #[test]
    fn test_sq_basic() {
        let dim = 128;
        let vectors = generate_random_vectors(1000, dim);
        
        let mut sq = ScalarQuantizer::for_dimension(dim);
        assert!(!sq.is_trained());
        
        sq.train(&vectors).unwrap();
        assert!(sq.is_trained());
        assert_eq!(sq.compression_ratio(), 4.0);
        assert_eq!(sq.code_size(), dim);
    }
    
    #[test]
    fn test_sq_quantize_dequantize() {
        let dim = 32;
        let vectors = generate_random_vectors(500, dim);
        
        let mut sq = ScalarQuantizer::for_dimension(dim);
        sq.train(&vectors).unwrap();
        
        // Quantize and dequantize
        let original = &vectors[0];
        let codes = sq.quantize(original).unwrap();
        assert_eq!(codes.len(), dim);
        
        let reconstructed = sq.dequantize(&codes).unwrap();
        assert_eq!(reconstructed.len(), dim);
        
        // Check reconstruction error is reasonable
        let error: f32 = original.iter()
            .zip(reconstructed.iter())
            .map(|(&a, &b)| (a - b).powi(2))
            .sum::<f32>()
            .sqrt();
        
        println!("SQ reconstruction error: {}", error);
        // Error should be small for 8-bit quantization
        assert!(error < 1.0, "Reconstruction error too high: {}", error);
    }
    
    #[test]
    fn test_sq_distance() {
        let dim = 64;
        let vectors = generate_random_vectors(500, dim);
        
        let mut sq = ScalarQuantizer::for_dimension(dim);
        sq.train(&vectors).unwrap();
        
        let query = &vectors[0];
        let codes = sq.quantize(&vectors[1]).unwrap();
        
        // Compute distance using SQ
        let sq_dist = sq.distance(query, &codes).unwrap();
        
        // Compute exact distance
        let exact_dist: f32 = query.iter()
            .zip(vectors[1].iter())
            .map(|(&a, &b)| (a - b).powi(2))
            .sum::<f32>()
            .sqrt();
        
        // SQ distance should be close to exact distance
        let relative_error = (sq_dist - exact_dist).abs() / exact_dist.max(1e-6);
        println!("SQ distance: {}, Exact: {}, Relative error: {:.2}%", 
                 sq_dist, exact_dist, relative_error * 100.0);
        
        assert!(relative_error < 0.1, "Distance error too high: {:.2}%", relative_error * 100.0);
    }
    
    #[test]
    fn test_sq_batch_quantize() {
        let dim = 32;
        let vectors = generate_random_vectors(100, dim);
        
        let mut sq = ScalarQuantizer::for_dimension(dim);
        sq.train(&vectors).unwrap();
        
        let codes = sq.quantize_batch(&vectors).unwrap();
        assert_eq!(codes.len(), 100);
        
        for code in codes {
            assert_eq!(code.len(), dim);
        }
    }
    
    #[test]
    fn test_sq_global_minmax() {
        let dim = 16;
        let vectors = generate_random_vectors(500, dim);
        
        let config = SQConfig { per_dimension: false };
        let mut sq = ScalarQuantizer::new(dim, config);
        sq.train(&vectors).unwrap();
        
        let codes = sq.quantize(&vectors[0]).unwrap();
        let reconstructed = sq.dequantize(&codes).unwrap();
        
        assert_eq!(reconstructed.len(), dim);
    }
    
    #[test]
    fn test_sq_errors() {
        let dim = 16;
        let sq = ScalarQuantizer::for_dimension(dim);
        
        // Not trained error
        let result = sq.quantize(&vec![0.0; dim]);
        assert!(matches!(result, Err(SQError::NotTrained)));
        
        // Dimension mismatch
        let vectors = generate_random_vectors(100, dim);
        let mut sq = ScalarQuantizer::for_dimension(dim);
        sq.train(&vectors).unwrap();
        
        let result = sq.quantize(&vec![0.0; 8]);
        assert!(matches!(result, Err(SQError::DimensionMismatch { .. })));
    }
    
    // =========================================================================
    // Product Quantization Tests
    // =========================================================================
    
    #[test]
    fn test_pq_config_for_dimension() {
        let config = PQConfig::for_dimension(1536);
        assert_eq!(1536 % config.num_subspaces, 0);
        
        let config = PQConfig::for_dimension(768);
        assert_eq!(768 % config.num_subspaces, 0);
        
        let config = PQConfig::for_dimension(128);
        assert_eq!(128 % config.num_subspaces, 0);
    }
    
    #[test]
    fn test_pq_train_and_encode() {
        let dim = 32;
        let vectors = generate_random_vectors(1000, dim);
        
        let config = PQConfig {
            num_subspaces: 4,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 500,
        };
        
        let mut pq = ProductQuantizer::new(dim, config);
        
        assert!(!pq.is_trained());
        
        pq.train(&vectors).unwrap();
        
        assert!(pq.is_trained());
        assert_eq!(pq.code_size(), 4);
        
        // Encode a vector
        let code = pq.encode(&vectors[0]).unwrap();
        assert_eq!(code.len(), 4);
        
        // Decode and check it's approximately the same
        let decoded = pq.decode(&code).unwrap();
        assert_eq!(decoded.len(), dim);
    }
    
    #[test]
    fn test_pq_compression_ratio() {
        let dim = 1536;
        let config = PQConfig {
            num_subspaces: 192,
            num_centroids: 256,
            training_iterations: 5,
            training_sample_size: 100,
        };
        
        let pq = ProductQuantizer::new(dim, config);
        
        // 1536 * 4 bytes = 6144 bytes original
        // 192 bytes compressed
        // Ratio = 6144 / 192 = 32
        let ratio = pq.compression_ratio();
        assert!((ratio - 32.0).abs() < 0.1);
    }
    
    #[test]
    fn test_pq_asymmetric_distance() {
        let dim = 16;
        let vectors = generate_random_vectors(500, dim);
        
        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 0,
        };
        
        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();
        
        // Encode a vector
        let code = pq.encode(&vectors[0]).unwrap();
        
        // Compute asymmetric distance
        let dist = pq.asymmetric_distance(&vectors[1], &code).unwrap();
        
        // Should be non-negative
        assert!(dist >= 0.0);
    }
    
    #[test]
    fn test_pq_distance_table() {
        let dim = 16;
        let vectors = generate_random_vectors(500, dim);
        
        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 0,
        };
        
        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();
        
        let query = &vectors[0];
        let code = pq.encode(&vectors[1]).unwrap();
        
        // Compute distance using table
        let table = pq.compute_distance_table(query).unwrap();
        let dist_table = pq.distance_with_table(&table, &code);
        
        // Compute distance directly
        let dist_direct = pq.asymmetric_distance(query, &code).unwrap();
        
        // Should be approximately equal
        assert!((dist_table - dist_direct).abs() < 1e-5);
    }
    
    #[test]
    fn test_pq_reconstruction_error() {
        let dim = 32;
        let vectors = generate_random_vectors(1000, dim);
        
        let config = PQConfig {
            num_subspaces: 4,
            num_centroids: 256,
            training_iterations: 20,
            training_sample_size: 0,
        };
        
        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();
        
        // Check reconstruction error
        let mut total_error = 0.0f32;
        
        for v in &vectors[..100] {
            let code = pq.encode(v).unwrap();
            let decoded = pq.decode(&code).unwrap();
            
            let error: f32 = v.iter()
                .zip(decoded.iter())
                .map(|(&a, &b)| (a - b).powi(2))
                .sum();
            
            total_error += error.sqrt();
        }
        
        let avg_error = total_error / 100.0;
        println!("Average reconstruction error: {}", avg_error);
        
        // Error should be reasonable (not exact, but not terrible)
        // The actual threshold depends on data distribution
        assert!(avg_error < 5.0, "Reconstruction error too high: {}", avg_error);
    }
    
    #[test]
    fn test_pq_batch_encode() {
        let dim = 16;
        let vectors = generate_random_vectors(100, dim);
        
        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 0,
        };
        
        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();
        
        let codes = pq.encode_batch(&vectors).unwrap();
        
        assert_eq!(codes.len(), 100);
        for code in codes {
            assert_eq!(code.len(), 2);
        }
    }
    
    #[test]
    fn test_pq_errors() {
        let dim = 16;
        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 0,
        };
        
        let pq = ProductQuantizer::new(dim, config);
        
        // Not trained error
        let result = pq.encode(&vec![0.0; dim]);
        assert!(matches!(result, Err(PQError::NotTrained)));
        
        // Dimension mismatch
        let mut pq = ProductQuantizer::new(dim, PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 5,
            training_sample_size: 0,
        });
        
        let vectors = generate_random_vectors(100, dim);
        pq.train(&vectors).unwrap();
        
        let result = pq.encode(&vec![0.0; 8]);
        assert!(matches!(result, Err(PQError::DimensionMismatch { .. })));
    }

    // =========================================================================
    // GPU fallback / DeviceSelector integration tests (Phase-63)
    // =========================================================================

    #[test]
    fn test_sq_train_with_cpu_selector() {
        // Explicitly force CPU selector — should use CPU path, produce same results
        let dim = 32;
        let vectors = generate_random_vectors(500, dim);

        let selector = crate::device::DeviceSelector::new(
            crate::device::GpuAccelConfig::cpu_only(),
        );

        let mut sq = ScalarQuantizer::for_dimension(dim);
        sq.train_with_selector(&vectors, Some(&selector)).unwrap();
        assert!(sq.is_trained());

        // Verify quantize/dequantize still works
        let codes = sq.quantize(&vectors[0]).unwrap();
        assert_eq!(codes.len(), dim);
        let recon = sq.dequantize(&codes).unwrap();
        assert_eq!(recon.len(), dim);
    }

    #[test]
    fn test_sq_train_with_auto_selector() {
        // Auto selector — will use GPU if available, otherwise CPU
        let dim = 16;
        let vectors = generate_random_vectors(100, dim);

        let selector = crate::device::DeviceSelector::auto();

        let mut sq = ScalarQuantizer::for_dimension(dim);
        sq.train_with_selector(&vectors, Some(&selector)).unwrap();
        assert!(sq.is_trained());

        let codes = sq.quantize(&vectors[0]).unwrap();
        assert_eq!(codes.len(), dim);
    }

    #[test]
    fn test_pq_train_with_cpu_selector() {
        // Force CPU — should produce valid codebook
        let dim = 16;
        let vectors = generate_random_vectors(500, dim);

        let selector = crate::device::DeviceSelector::new(
            crate::device::GpuAccelConfig::cpu_only(),
        );

        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 0,
        };

        let mut pq = ProductQuantizer::new(dim, config);
        pq.train_with_selector(&vectors, Some(&selector)).unwrap();
        assert!(pq.is_trained());

        let code = pq.encode(&vectors[0]).unwrap();
        assert_eq!(code.len(), 2);
        let decoded = pq.decode(&code).unwrap();
        assert_eq!(decoded.len(), dim);
    }

    #[test]
    fn test_pq_train_with_auto_selector() {
        // Auto selector — GPU or CPU
        let dim = 16;
        let vectors = generate_random_vectors(200, dim);

        let selector = crate::device::DeviceSelector::auto();

        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 8,
            training_iterations: 5,
            training_sample_size: 0,
        };

        let mut pq = ProductQuantizer::new(dim, config);
        pq.train_with_selector(&vectors, Some(&selector)).unwrap();
        assert!(pq.is_trained());

        // Verify encode still works
        let code = pq.encode(&vectors[0]).unwrap();
        assert_eq!(code.len(), 2);
    }

    #[test]
    fn test_pq_cpu_gpu_consistency() {
        // Train with CPU-only selector and auto selector (on CPU machine,
        // both should produce similar results structurally)
        let dim = 16;
        let vectors = generate_random_vectors(200, dim);

        let cpu_selector = crate::device::DeviceSelector::new(
            crate::device::GpuAccelConfig::cpu_only(),
        );

        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 8,
            training_iterations: 10,
            training_sample_size: 0,
        };

        let mut pq_cpu = ProductQuantizer::new(dim, config.clone());
        pq_cpu.train_with_selector(&vectors, Some(&cpu_selector)).unwrap();

        // Both should be trained and produce valid codes
        assert!(pq_cpu.is_trained());
        let code = pq_cpu.encode(&vectors[0]).unwrap();
        assert_eq!(code.len(), 2);

        // Verify compression ratio is correct
        let ratio = pq_cpu.compression_ratio();
        assert!(ratio > 1.0, "Compression ratio should be > 1: {ratio}");
    }

    // =========================================================================
    // PQ distance table GPU tests (Phase-67)
    // =========================================================================

    #[test]
    fn test_pq_distance_table_gpu_cpu_consistency() {
        // GPU and CPU distance table should produce identical results
        let dim = 16;
        let vectors = generate_random_vectors(500, dim);

        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 10,
            training_sample_size: 0,
        };

        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();

        let query = &vectors[0];

        // CPU table
        let cpu_table = pq.compute_distance_table_cpu(query).unwrap();

        // Full table (may use GPU if available, or CPU)
        let auto_table = pq.compute_distance_table(query).unwrap();

        // Must match in dimensions
        assert_eq!(cpu_table.len(), auto_table.len());
        for (s, (ct, at)) in cpu_table.iter().zip(auto_table.iter()).enumerate() {
            assert_eq!(ct.len(), at.len(), "subspace {} centroid count mismatch", s);
            for (i, (c, a)) in ct.iter().zip(at.iter()).enumerate() {
                assert!(
                    (c - a).abs() < 1e-2,
                    "subspace {} centroid {} distance mismatch: cpu={} auto={}",
                    s, i, c, a
                );
            }
        }
    }

    #[test]
    fn test_pq_distance_table_with_search() {
        // End-to-end: train PQ, compute distance table, verify search ranking
        let dim = 32;
        let vectors = generate_random_vectors(1000, dim);

        let config = PQConfig {
            num_subspaces: 4,
            num_centroids: 256,
            training_iterations: 15,
            training_sample_size: 0,
        };

        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();

        let query = &vectors[0];
        let table = pq.compute_distance_table(query).unwrap();

        // Encode all vectors and compute distances using table
        let mut table_dists: Vec<(usize, f32)> = Vec::new();
        let mut direct_dists: Vec<(usize, f32)> = Vec::new();

        for (i, v) in vectors.iter().enumerate().skip(1) {
            let code = pq.encode(v).unwrap();
            let td = pq.distance_with_table(&table, &code);
            let dd = pq.asymmetric_distance(query, &code).unwrap();
            table_dists.push((i, td));
            direct_dists.push((i, dd));
        }

        table_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        direct_dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

        // Top-10 from table and direct should match exactly
        let top10_table: Vec<usize> = table_dists.iter().take(10).map(|(i, _)| *i).collect();
        let top10_direct: Vec<usize> = direct_dists.iter().take(10).map(|(i, _)| *i).collect();
        assert_eq!(top10_table, top10_direct, "Table vs direct top-10 should match");
    }

    #[test]
    fn test_pq_distance_table_small_centroids() {
        // Test with fewer centroids (8 = 3-bit codes)
        let dim = 8;
        let vectors = generate_random_vectors(200, dim);

        let config = PQConfig {
            num_subspaces: 2,
            num_centroids: 8,
            training_iterations: 10,
            training_sample_size: 0,
        };

        let mut pq = ProductQuantizer::new(dim, config);
        pq.train(&vectors).unwrap();

        let query = &vectors[0];
        let table = pq.compute_distance_table(query).unwrap();

        assert_eq!(table.len(), 2); // 2 subspaces
        assert_eq!(table[0].len(), 8); // 8 centroids
        assert_eq!(table[1].len(), 8);

        // All distances should be non-negative
        for sub_dists in &table {
            for &d in sub_dists {
                assert!(d >= 0.0, "Distance should be non-negative: {}", d);
            }
        }
    }
}
