//! GPU-accelerated vector quantization training (PQ + SQ)
//!
//! Implements Product Quantization (PQ) and Scalar Quantization (SQ) training
//! on GPU for IVF index construction.
//!
//! ## Product Quantization (PQ)
//! Splits each vector into sub-vectors and trains a codebook per sub-space.
//! GPU accelerates the per-subspace KMeans training.
//!
//! ## Scalar Quantization (SQ)
//! Computes global min/max per dimension, then linearly maps to uint8.
//! GPU accelerates the min/max reduction scan.

use super::error::{GpuError, GpuResult};
use super::kmeans::{GpuKMeansConfig, GpuKMeansTrainer, KMeansResult};
use super::memory::GpuMemoryPool;
use crate::device::{DeviceSelector, WorkloadHint, GpuFallbackEvent, FallbackReason};

// ============================================================================
// Scalar Quantization (SQ)
// ============================================================================

/// Configuration for Scalar Quantization training.
#[derive(Debug, Clone)]
pub struct SqTrainConfig {
    /// Number of bits (4 or 8)
    pub num_bits: u8,
    /// Sample rate for training (how many vectors to sample for min/max)
    pub sample_rate: usize,
}

impl Default for SqTrainConfig {
    fn default() -> Self {
        Self {
            num_bits: 8,
            sample_rate: 256,
        }
    }
}

/// Result of SQ training — the quantization parameters.
#[derive(Debug, Clone)]
pub struct SqParams {
    /// Per-dimension minimum values [dimension]
    pub min_vals: Vec<f32>,
    /// Per-dimension maximum values [dimension]
    pub max_vals: Vec<f32>,
    /// Per-dimension scale factors: (max - min) / (2^bits - 1)
    pub scales: Vec<f32>,
    /// Number of quantization bits
    pub num_bits: u8,
    /// Vector dimension
    pub dimension: usize,
    /// Whether GPU was used for training
    pub used_gpu: bool,
    /// Training time in milliseconds
    pub train_time_ms: u64,
}

impl SqParams {
    /// Quantize a single f32 vector to uint8 codes.
    pub fn quantize(&self, vector: &[f32]) -> Vec<u8> {
        let max_val = (1u32 << self.num_bits) - 1;
        vector
            .iter()
            .enumerate()
            .map(|(d, &v)| {
                if self.scales[d] == 0.0 {
                    0u8
                } else {
                    let normalized = (v - self.min_vals[d]) / self.scales[d];
                    let clamped = normalized.clamp(0.0, max_val as f32);
                    clamped.round() as u8
                }
            })
            .collect()
    }

    /// Dequantize uint8 codes back to f32 vector (approximate).
    pub fn dequantize(&self, codes: &[u8]) -> Vec<f32> {
        codes
            .iter()
            .enumerate()
            .map(|(d, &code)| {
                self.min_vals[d] + code as f32 * self.scales[d]
            })
            .collect()
    }
}

/// GPU-accelerated Scalar Quantization trainer.
pub struct GpuSqTrainer {
    config: SqTrainConfig,
}

impl GpuSqTrainer {
    pub fn new(config: SqTrainConfig) -> Self {
        Self { config }
    }

    /// Train SQ parameters from vectors.
    ///
    /// Computes per-dimension min/max across all vectors.
    /// Uses GPU for parallel reduction if available.
    pub fn train(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> GpuResult<SqParams> {
        self.train_with_selector(vectors, n_vectors, dimension, memory_pool, None)
    }

    /// Train SQ with explicit DeviceSelector for unified GPU device management.
    pub fn train_with_selector(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
        selector: Option<&DeviceSelector>,
    ) -> GpuResult<SqParams> {
        if n_vectors == 0 || dimension == 0 {
            return Err(GpuError::Other("Empty input for SQ training".into()));
        }

        let start_time = std::time::Instant::now();

        // Use DeviceSelector for intelligent GPU decision
        let sel = selector.cloned().unwrap_or_else(DeviceSelector::auto);
        let hint = WorkloadHint::with_operation(n_vectors, dimension, "gpu_sq_train");
        let use_gpu = sel.should_use_gpu(&hint);

        // Try GPU path
        #[cfg(feature = "candle")]
        if use_gpu {
            match self.train_gpu_with_selector(vectors, n_vectors, dimension, memory_pool, &sel) {
                Ok(params) => return Ok(params),
                Err(e) => {
                    tracing::warn!("GPU SQ training failed, falling back to CPU: {}", e);
                    sel.emit_fallback_public(GpuFallbackEvent {
                        reason: FallbackReason::RuntimeError,
                        detail: format!("GPU SQ train error: {e}"),
                        operation: "gpu_sq_train".into(),
                        requested_device: sel.config().device_preference.clone(),
                    });
                }
            }
        }

        // CPU fallback
        let params = self.train_cpu(vectors, n_vectors, dimension);
        let elapsed = start_time.elapsed().as_millis() as u64;

        Ok(SqParams {
            train_time_ms: elapsed,
            ..params
        })
    }

    /// GPU-accelerated SQ training (via DeviceSelector).
    #[cfg(feature = "candle")]
    fn train_gpu_with_selector(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        _memory_pool: Option<&GpuMemoryPool>,
        selector: &DeviceSelector,
    ) -> GpuResult<SqParams> {
        use candle_core::{Device, Tensor};

        let start_time = std::time::Instant::now();

        let (device, is_gpu) = selector.candle_device_with_fallback("gpu_sq_train");
        if !is_gpu {
            return Err(GpuError::InitializationError(
                "DeviceSelector returned CPU; GPU not available".into(),
            ));
        }

        // Transfer to GPU
        let tensor = Tensor::from_slice(vectors, (n_vectors, dimension), &device)
            .map_err(|e| GpuError::TransferError(format!("upload vectors: {e}")))?;

        // Compute per-column min and max using GPU reduction
        let min_tensor = tensor.min(0)
            .map_err(|e| GpuError::KernelError(format!("min reduction: {e}")))?;
        let max_tensor = tensor.max(0)
            .map_err(|e| GpuError::KernelError(format!("max reduction: {e}")))?;

        // Transfer results back to CPU
        let min_vals: Vec<f32> = min_tensor
            .to_device(&Device::Cpu)
            .map_err(|e| GpuError::TransferError(format!("copy min: {e}")))?
            .to_vec1::<f32>()
            .map_err(|e| GpuError::TransferError(format!("convert min: {e}")))?;

        let max_vals: Vec<f32> = max_tensor
            .to_device(&Device::Cpu)
            .map_err(|e| GpuError::TransferError(format!("copy max: {e}")))?
            .to_vec1::<f32>()
            .map_err(|e| GpuError::TransferError(format!("convert max: {e}")))?;

        let max_code = (1u32 << self.config.num_bits) - 1;
        let scales: Vec<f32> = min_vals
            .iter()
            .zip(max_vals.iter())
            .map(|(min, max)| {
                let range = max - min;
                if range == 0.0 { 0.0 } else { range / max_code as f32 }
            })
            .collect();

        let elapsed = start_time.elapsed().as_millis() as u64;

        tracing::info!(
            "GPU SQ training completed: {}d, {} bits, {:.1}ms",
            dimension, self.config.num_bits, elapsed
        );

        Ok(SqParams {
            min_vals,
            max_vals,
            scales,
            num_bits: self.config.num_bits,
            dimension,
            used_gpu: true,
            train_time_ms: elapsed,
        })
    }

    /// CPU fallback SQ training.
    fn train_cpu(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
    ) -> SqParams {
        let start_time = std::time::Instant::now();

        let mut min_vals = vec![f32::MAX; dimension];
        let mut max_vals = vec![f32::MIN; dimension];

        for i in 0..n_vectors {
            let v = &vectors[i * dimension..(i + 1) * dimension];
            for (d, &val) in v.iter().enumerate() {
                min_vals[d] = min_vals[d].min(val);
                max_vals[d] = max_vals[d].max(val);
            }
        }

        let max_code = (1u32 << self.config.num_bits) - 1;
        let scales: Vec<f32> = min_vals
            .iter()
            .zip(max_vals.iter())
            .map(|(min, max)| {
                let range = max - min;
                if range == 0.0 { 0.0 } else { range / max_code as f32 }
            })
            .collect();

        let elapsed = start_time.elapsed().as_millis() as u64;

        SqParams {
            min_vals,
            max_vals,
            scales,
            num_bits: self.config.num_bits,
            dimension,
            used_gpu: false,
            train_time_ms: elapsed,
        }
    }
}

// ============================================================================
// Product Quantization (PQ)
// ============================================================================

/// Configuration for Product Quantization training.
#[derive(Debug, Clone)]
pub struct PqTrainConfig {
    /// Number of sub-vectors to split each vector into.
    /// Dimension must be divisible by this. Typical: dim / 8.
    pub num_sub_vectors: usize,

    /// Number of centroids per sub-codebook (default: 256 for 8-bit PQ).
    pub num_centroids: usize,

    /// Number of bits per centroid code (default: 8).
    pub num_bits: u8,

    /// KMeans config for sub-codebook training.
    pub kmeans_config: GpuKMeansConfig,
}

impl Default for PqTrainConfig {
    fn default() -> Self {
        Self {
            num_sub_vectors: 8,
            num_centroids: 256,
            num_bits: 8,
            kmeans_config: GpuKMeansConfig {
                n_clusters: 256,
                max_iterations: 50,
                tolerance: 1e-4,
                seed: 42,
                ..Default::default()
            },
        }
    }
}

/// Result of PQ training — the codebooks.
#[derive(Debug, Clone)]
pub struct PqCodebook {
    /// Codebook data: [num_sub_vectors][num_centroids][sub_dimension]
    /// Flattened as: [sub_vec_idx * num_centroids * sub_dim + centroid_idx * sub_dim + d]
    pub codebooks: Vec<f32>,

    /// Number of sub-vectors
    pub num_sub_vectors: usize,

    /// Number of centroids per sub-codebook
    pub num_centroids: usize,

    /// Dimension of each sub-vector
    pub sub_dimension: usize,

    /// Full vector dimension
    pub dimension: usize,

    /// Per sub-codebook training results
    pub sub_results: Vec<SubCodebookResult>,

    /// Whether GPU was used
    pub used_gpu: bool,

    /// Total training time in milliseconds
    pub train_time_ms: u64,
}

/// Training result for a single sub-codebook.
#[derive(Debug, Clone)]
pub struct SubCodebookResult {
    /// Sub-vector index
    pub sub_index: usize,
    /// Number of KMeans iterations
    pub iterations: usize,
    /// Final inertia
    pub inertia: f64,
    /// Whether it converged
    pub converged: bool,
}

impl PqCodebook {
    /// Encode a single vector to PQ codes.
    pub fn encode(&self, vector: &[f32]) -> Vec<u8> {
        let mut codes = Vec::with_capacity(self.num_sub_vectors);

        for s in 0..self.num_sub_vectors {
            let sub_start = s * self.sub_dimension;
            let sub_end = sub_start + self.sub_dimension;
            let sub_vec = &vector[sub_start..sub_end];

            // Find nearest centroid in this sub-codebook
            let codebook_offset = s * self.num_centroids * self.sub_dimension;
            let mut best_code = 0u8;
            let mut best_dist = f32::MAX;

            for c in 0..self.num_centroids {
                let cent_start = codebook_offset + c * self.sub_dimension;
                let centroid = &self.codebooks[cent_start..cent_start + self.sub_dimension];

                let dist: f32 = sub_vec
                    .iter()
                    .zip(centroid)
                    .map(|(a, b)| (a - b) * (a - b))
                    .sum();

                if dist < best_dist {
                    best_dist = dist;
                    best_code = c as u8;
                }
            }

            codes.push(best_code);
        }

        codes
    }

    /// Decode PQ codes back to approximate vector.
    pub fn decode(&self, codes: &[u8]) -> Vec<f32> {
        let mut vector = Vec::with_capacity(self.dimension);

        for (s, &code) in codes.iter().enumerate() {
            let codebook_offset = s * self.num_centroids * self.sub_dimension;
            let cent_start = codebook_offset + code as usize * self.sub_dimension;
            let centroid = &self.codebooks[cent_start..cent_start + self.sub_dimension];
            vector.extend_from_slice(centroid);
        }

        vector
    }

    /// Get compression ratio (original bytes / compressed bytes).
    pub fn compression_ratio(&self) -> f64 {
        let original = self.dimension * std::mem::size_of::<f32>();
        let compressed = self.num_sub_vectors; // 1 byte per sub-vector (for 8-bit PQ)
        original as f64 / compressed as f64
    }
}

/// GPU-accelerated Product Quantization trainer.
pub struct GpuPqTrainer {
    config: PqTrainConfig,
}

impl GpuPqTrainer {
    pub fn new(config: PqTrainConfig) -> Self {
        Self { config }
    }

    /// Train PQ codebooks from vectors.
    ///
    /// Splits vectors into sub-vectors and trains a codebook per sub-space
    /// using KMeans (GPU-accelerated if available).
    pub fn train(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> GpuResult<PqCodebook> {
        self.train_with_selector(vectors, n_vectors, dimension, memory_pool, None)
    }

    /// Train PQ codebooks with explicit DeviceSelector for unified GPU device management.
    pub fn train_with_selector(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
        selector: Option<&DeviceSelector>,
    ) -> GpuResult<PqCodebook> {
        if dimension % self.config.num_sub_vectors != 0 {
            return Err(GpuError::Other(format!(
                "Dimension {} not divisible by num_sub_vectors {}",
                dimension, self.config.num_sub_vectors
            )));
        }

        let sub_dimension = dimension / self.config.num_sub_vectors;
        let start_time = std::time::Instant::now();

        let mut all_codebooks = Vec::with_capacity(
            self.config.num_sub_vectors * self.config.num_centroids * sub_dimension,
        );
        let mut sub_results = Vec::with_capacity(self.config.num_sub_vectors);
        let mut any_used_gpu = false;

        // Train each sub-codebook
        for s in 0..self.config.num_sub_vectors {
            // Extract sub-vectors for this sub-space
            let mut sub_vectors = Vec::with_capacity(n_vectors * sub_dimension);
            for i in 0..n_vectors {
                let start = i * dimension + s * sub_dimension;
                sub_vectors.extend_from_slice(&vectors[start..start + sub_dimension]);
            }

            // Train KMeans on sub-vectors (pass DeviceSelector through)
            let kmeans_config = GpuKMeansConfig {
                n_clusters: self.config.num_centroids,
                seed: self.config.kmeans_config.seed + s as u64,
                ..self.config.kmeans_config.clone()
            };

            let trainer = GpuKMeansTrainer::new(kmeans_config);
            let result = trainer.train_with_selector(
                &sub_vectors, n_vectors, sub_dimension, memory_pool, selector,
            )?;

            if result.used_gpu {
                any_used_gpu = true;
            }

            all_codebooks.extend_from_slice(&result.centroids);
            sub_results.push(SubCodebookResult {
                sub_index: s,
                iterations: result.iterations,
                inertia: result.inertia,
                converged: result.converged,
            });

            if s % 4 == 0 || s == self.config.num_sub_vectors - 1 {
                tracing::debug!(
                    "PQ sub-codebook {}/{}: {} iterations, inertia={:.2}",
                    s + 1,
                    self.config.num_sub_vectors,
                    result.iterations,
                    result.inertia,
                );
            }
        }

        let elapsed = start_time.elapsed().as_millis() as u64;
        let converged_count = sub_results.iter().filter(|r| r.converged).count();

        tracing::info!(
            "PQ training completed: {}d → {}×{}d, {} centroids, {}/{} converged, {:.1}ms{}",
            dimension,
            self.config.num_sub_vectors,
            sub_dimension,
            self.config.num_centroids,
            converged_count,
            self.config.num_sub_vectors,
            elapsed,
            if any_used_gpu { " (GPU)" } else { " (CPU)" }
        );

        Ok(PqCodebook {
            codebooks: all_codebooks,
            num_sub_vectors: self.config.num_sub_vectors,
            num_centroids: self.config.num_centroids,
            sub_dimension,
            dimension,
            sub_results,
            used_gpu: any_used_gpu,
            train_time_ms: elapsed,
        })
    }
}

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

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

    fn generate_test_vectors(n: usize, dim: usize) -> Vec<f32> {
        (0..n * dim)
            .map(|i| ((i * 7 + 13) % 1000) as f32 / 100.0)
            .collect()
    }

    // ---- SQ Tests ----

    #[test]
    fn test_sq_train_cpu() {
        let data = generate_test_vectors(100, 8);
        let trainer = GpuSqTrainer::new(SqTrainConfig::default());
        let params = trainer.train_cpu(&data, 100, 8);

        assert_eq!(params.dimension, 8);
        assert_eq!(params.min_vals.len(), 8);
        assert_eq!(params.max_vals.len(), 8);
        assert_eq!(params.scales.len(), 8);
        assert!(!params.used_gpu);

        // Min should be <= max for each dimension
        for d in 0..8 {
            assert!(params.min_vals[d] <= params.max_vals[d]);
        }
    }

    #[test]
    fn test_sq_quantize_dequantize() {
        let data = vec![
            0.0, 1.0, 2.0, 3.0,
            4.0, 5.0, 6.0, 7.0,
            8.0, 9.0, 10.0, 11.0,
        ];
        let trainer = GpuSqTrainer::new(SqTrainConfig::default());
        let params = trainer.train_cpu(&data, 3, 4);

        // Quantize and dequantize
        let original = &[4.0f32, 5.0, 6.0, 7.0];
        let codes = params.quantize(original);
        assert_eq!(codes.len(), 4);

        let reconstructed = params.dequantize(&codes);
        assert_eq!(reconstructed.len(), 4);

        // Reconstruction error should be small
        for (o, r) in original.iter().zip(reconstructed.iter()) {
            assert!((o - r).abs() < 0.5, "Reconstruction error too large: {} vs {}", o, r);
        }
    }

    #[test]
    fn test_sq_edge_case_constant_dimension() {
        // All values the same in a dimension → scale = 0
        let data = vec![
            5.0, 1.0,
            5.0, 2.0,
            5.0, 3.0,
        ];
        let trainer = GpuSqTrainer::new(SqTrainConfig::default());
        let params = trainer.train_cpu(&data, 3, 2);

        assert_eq!(params.scales[0], 0.0); // Constant dimension
        assert!(params.scales[1] > 0.0);   // Varying dimension

        let codes = params.quantize(&[5.0, 2.0]);
        assert_eq!(codes[0], 0); // Constant → 0
    }

    // ---- PQ Tests ----

    #[test]
    fn test_pq_train_cpu() {
        let dim = 16;
        let n_sub = 4; // 16d / 4 = 4d sub-vectors
        let data = generate_test_vectors(200, dim);

        let config = PqTrainConfig {
            num_sub_vectors: n_sub,
            num_centroids: 8, // Small for testing
            kmeans_config: GpuKMeansConfig {
                n_clusters: 8,
                max_iterations: 10,
                seed: 42,
                ..Default::default()
            },
            ..Default::default()
        };

        let trainer = GpuPqTrainer::new(config);
        let codebook = trainer.train(&data, 200, dim, None).unwrap();

        assert_eq!(codebook.dimension, dim);
        assert_eq!(codebook.num_sub_vectors, n_sub);
        assert_eq!(codebook.sub_dimension, dim / n_sub);
        assert_eq!(codebook.num_centroids, 8);
        assert_eq!(codebook.sub_results.len(), n_sub);
        assert_eq!(
            codebook.codebooks.len(),
            n_sub * 8 * (dim / n_sub)
        );
    }

    #[test]
    fn test_pq_encode_decode() {
        let dim = 8;
        let data = generate_test_vectors(100, dim);

        let config = PqTrainConfig {
            num_sub_vectors: 2,
            num_centroids: 4,
            kmeans_config: GpuKMeansConfig {
                n_clusters: 4,
                max_iterations: 20,
                seed: 42,
                ..Default::default()
            },
            ..Default::default()
        };

        let trainer = GpuPqTrainer::new(config);
        let codebook = trainer.train(&data, 100, dim, None).unwrap();

        // Encode and decode a vector
        let original = &data[0..dim];
        let codes = codebook.encode(original);
        assert_eq!(codes.len(), 2); // 2 sub-vectors

        let decoded = codebook.decode(&codes);
        assert_eq!(decoded.len(), dim);

        // Compression ratio for 8d with 2 sub-vectors:
        // original = 8 * 4 = 32 bytes, compressed = 2 bytes
        let ratio = codebook.compression_ratio();
        assert!(ratio > 1.0, "Compression ratio should be > 1: {ratio}");
    }

    #[test]
    fn test_pq_dimension_mismatch() {
        let config = PqTrainConfig {
            num_sub_vectors: 3, // 10 is not divisible by 3
            ..Default::default()
        };

        let trainer = GpuPqTrainer::new(config);
        let data = generate_test_vectors(100, 10);
        let result = trainer.train(&data, 100, 10, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_sq_full_train_with_fallback() {
        let data = generate_test_vectors(50, 4);
        let trainer = GpuSqTrainer::new(SqTrainConfig::default());
        let result = trainer.train(&data, 50, 4, None).unwrap();
        assert_eq!(result.dimension, 4);
    }
}
