//! GPU-accelerated KMeans clustering for IVF index training
//!
//! Implements KMeans++ initialization and iterative Lloyd's algorithm
//! using GPU matrix operations via candle CUDA backend.
//!
//! ## Fallback Strategy
//!
//! If GPU is unavailable or fails mid-computation, each operation
//! transparently falls back to CPU (rayon-parallelized).
//!
//! ## Algorithm Overview
//!
//! 1. **KMeans++ initialization**: Select initial centroids with distance-weighted
//!    probability (runs on GPU for distance computation).
//! 2. **Assignment step**: Compute distances from all vectors to all centroids
//!    using cuBLAS-backed matrix multiplication (`||a-b||² = ||a||² + ||b||² - 2aᵀb`).
//! 3. **Update step**: Recompute centroids as mean of assigned vectors.
//! 4. **Convergence check**: Stop when inertia change < tolerance.

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

/// Configuration for GPU KMeans training.
#[derive(Debug, Clone)]
pub struct GpuKMeansConfig {
    /// Number of clusters (K)
    pub n_clusters: usize,

    /// Maximum number of iterations
    pub max_iterations: usize,

    /// Convergence tolerance (relative inertia change)
    pub tolerance: f64,

    /// Random seed for initialization
    pub seed: u64,

    /// Batch size for distance computation (0 = auto based on GPU memory)
    pub batch_size: usize,

    /// Whether to use KMeans++ initialization (recommended) or random
    pub use_kmeans_plus_plus: bool,

    /// Enable verbose logging of training progress
    pub verbose: bool,
}

impl Default for GpuKMeansConfig {
    fn default() -> Self {
        Self {
            n_clusters: 256,
            max_iterations: 50,
            tolerance: 1e-4,
            seed: 42,
            batch_size: 0,
            use_kmeans_plus_plus: true,
            verbose: false,
        }
    }
}

/// Result of KMeans training.
#[derive(Debug, Clone)]
pub struct KMeansResult {
    /// Final centroids matrix [n_clusters × dimension]
    pub centroids: Vec<f32>,

    /// Number of clusters
    pub n_clusters: usize,

    /// Vector dimension
    pub dimension: usize,

    /// Cluster assignment for each input vector
    pub assignments: Vec<u32>,

    /// Number of vectors in each cluster
    pub cluster_sizes: Vec<usize>,

    /// Final inertia (sum of squared distances to nearest centroid)
    pub inertia: f64,

    /// Number of iterations performed
    pub iterations: usize,

    /// Whether training converged
    pub converged: bool,

    /// Whether GPU was used (false = fell back to CPU)
    pub used_gpu: bool,

    /// Training time in milliseconds
    pub train_time_ms: u64,
}

impl KMeansResult {
    /// Get centroid for a specific cluster.
    pub fn centroid(&self, cluster_id: usize) -> Option<&[f32]> {
        if cluster_id >= self.n_clusters {
            return None;
        }
        let start = cluster_id * self.dimension;
        let end = start + self.dimension;
        Some(&self.centroids[start..end])
    }
}

/// GPU-accelerated KMeans trainer.
///
/// Attempts to use GPU for distance computation and centroid updates.
/// Falls back to CPU if GPU is unavailable or encounters errors.
pub struct GpuKMeansTrainer {
    config: GpuKMeansConfig,
}

impl GpuKMeansTrainer {
    /// Create a new GPU KMeans trainer.
    pub fn new(config: GpuKMeansConfig) -> Self {
        Self { config }
    }

    /// Train KMeans on the given vectors.
    ///
    /// `vectors` is a flat array of shape [n_vectors × dimension].
    /// If GPU is available, uses GPU-accelerated distance computation.
    /// Otherwise, falls back to CPU automatically.
    pub fn train(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> GpuResult<KMeansResult> {
        self.train_with_selector(vectors, n_vectors, dimension, memory_pool, None)
    }

    /// Train KMeans with explicit DeviceSelector for unified GPU device management.
    ///
    /// When `selector` is `None`, falls back to legacy auto-detection behavior.
    /// When provided, uses `DeviceSelector::should_use_gpu()` for intelligent decision.
    pub fn train_with_selector(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
        selector: Option<&DeviceSelector>,
    ) -> GpuResult<KMeansResult> {
        if n_vectors == 0 || dimension == 0 {
            return Err(GpuError::Other("Empty input vectors".into()));
        }
        if vectors.len() != n_vectors * dimension {
            return Err(GpuError::Other(format!(
                "Vector data size mismatch: expected {}, got {}",
                n_vectors * dimension,
                vectors.len()
            )));
        }
        if self.config.n_clusters > n_vectors {
            return Err(GpuError::Other(format!(
                "n_clusters ({}) > n_vectors ({})",
                self.config.n_clusters, n_vectors
            )));
        }

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

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

        // Try GPU path if DeviceSelector says yes
        #[cfg(feature = "candle")]
        if use_gpu {
            let sel = selector.cloned().unwrap_or_else(DeviceSelector::auto);
            match self.train_gpu_with_selector(vectors, n_vectors, dimension, memory_pool, &sel) {
                Ok(result) => return Ok(result),
                Err(e) => {
                    tracing::warn!(
                        "GPU KMeans failed, falling back to CPU: {}",
                        e
                    );
                    // Log fallback event
                    sel.emit_fallback_public(GpuFallbackEvent {
                        reason: FallbackReason::RuntimeError,
                        detail: format!("GPU KMeans train error: {e}"),
                        operation: "gpu_kmeans".into(),
                        requested_device: sel.config().device_preference.clone(),
                    });
                    // Fall through to CPU path
                }
            }
        }

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

        Ok(KMeansResult {
            train_time_ms: elapsed,
            ..result
        })
    }

    /// GPU-accelerated KMeans training using candle CUDA (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<KMeansResult> {
        use candle_core::{Device, Tensor, DType};

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

        // Use DeviceSelector to get device (auto fallback to CPU is handled by caller)
        let (device, is_gpu) = selector.candle_device_with_fallback("gpu_kmeans_train");
        if !is_gpu {
            return Err(GpuError::InitializationError(
                "DeviceSelector returned CPU device; GPU not available".into(),
            ));
        }

        // Check memory budget if pool is available
        let data_size = n_vectors * dimension * std::mem::size_of::<f32>();
        let centroid_size = self.config.n_clusters * dimension * std::mem::size_of::<f32>();
        let working_memory = data_size + centroid_size * 3; // data + centroids + distances + assignments

        if let Some(pool) = memory_pool {
            let reservation = pool.try_reserve(working_memory);
            if reservation.is_cpu_spill() {
                return Err(GpuError::OutOfMemory {
                    requested_bytes: working_memory,
                    available_bytes: pool.budget().saturating_sub(pool.allocated()),
                    device_name: "GPU".into(),
                });
            }
        }

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

        // Initialize centroids
        let mut centroids = if self.config.use_kmeans_plus_plus {
            self.kmeans_plus_plus_init_gpu(&vectors_tensor, n_vectors, dimension, &device)?
        } else {
            self.random_init_gpu(&vectors_tensor, dimension, &device)?
        };

        // Precompute vector norms: ||x||² for each vector
        let vector_norms = vectors_tensor
            .sqr().map_err(|e| GpuError::KernelError(format!("sqr failed: {e}")))?
            .sum(1).map_err(|e| GpuError::KernelError(format!("sum failed: {e}")))?;

        let mut prev_inertia = f64::MAX;
        let mut iterations = 0;
        let mut converged = false;
        let mut assignments = vec![0u32; n_vectors];
        let mut cluster_sizes = vec![0usize; self.config.n_clusters];

        for iter in 0..self.config.max_iterations {
            iterations = iter + 1;

            // ---- Assignment step ----
            // Compute distances using: ||a-b||² = ||a||² + ||b||² - 2aᵀb
            let centroid_norms = centroids
                .sqr().map_err(|e| GpuError::KernelError(format!("centroid sqr: {e}")))?
                .sum(1).map_err(|e| GpuError::KernelError(format!("centroid sum: {e}")))?;

            // -2 * vectors @ centroids.T → [n_vectors, n_clusters]
            let neg2_dot = vectors_tensor
                .matmul(&centroids.t().map_err(|e| GpuError::KernelError(format!("transpose: {e}")))?)
                .map_err(|e| GpuError::KernelError(format!("matmul: {e}")))?
                .affine(-2.0, 0.0)
                .map_err(|e| GpuError::KernelError(format!("affine: {e}")))?;

            // distances = ||a||² + ||b||² - 2aᵀb
            let vector_norms_2d = vector_norms
                .unsqueeze(1)
                .map_err(|e| GpuError::KernelError(format!("unsqueeze: {e}")))?;
            let centroid_norms_2d = centroid_norms
                .unsqueeze(0)
                .map_err(|e| GpuError::KernelError(format!("unsqueeze: {e}")))?;

            let distances = neg2_dot
                .broadcast_add(&vector_norms_2d)
                .map_err(|e| GpuError::KernelError(format!("broadcast_add: {e}")))?
                .broadcast_add(&centroid_norms_2d)
                .map_err(|e| GpuError::KernelError(format!("broadcast_add: {e}")))?;

            // argmin per row → assignments [n_vectors]
            let assignments_tensor = distances
                .argmin(1)
                .map_err(|e| GpuError::KernelError(format!("argmin: {e}")))?;

            // min per row → inertia
            let min_distances = distances
                .min(1)
                .map_err(|e| GpuError::KernelError(format!("min: {e}")))?;
            let inertia_tensor = min_distances
                .sum_all()
                .map_err(|e| GpuError::KernelError(format!("sum_all: {e}")))?;

            // Copy back to CPU for update step
            let assignments_cpu = assignments_tensor
                .to_device(&Device::Cpu)
                .map_err(|e| GpuError::TransferError(format!("copy assignments: {e}")))?;
            let assign_vec: Vec<u32> = assignments_cpu
                .to_vec1::<u32>()
                .map_err(|e| GpuError::TransferError(format!("convert assignments: {e}")))?;

            let inertia: f64 = inertia_tensor
                .to_device(&Device::Cpu)
                .map_err(|e| GpuError::TransferError(format!("copy inertia: {e}")))?
                .to_scalar::<f32>()
                .map_err(|e| GpuError::TransferError(format!("scalar inertia: {e}")))? as f64;

            assignments = assign_vec;

            // Check convergence
            let inertia_change = ((prev_inertia - inertia) / prev_inertia).abs();
            if self.config.verbose {
                tracing::info!(
                    "GPU KMeans iter {}/{}: inertia={:.2}, change={:.6}",
                    iterations,
                    self.config.max_iterations,
                    inertia,
                    inertia_change
                );
            }

            if inertia_change < self.config.tolerance && iter > 0 {
                converged = true;
                prev_inertia = inertia;
                break;
            }
            prev_inertia = inertia;

            // ---- Update step (CPU, then transfer back to GPU) ----
            // Recompute centroids as mean of assigned vectors
            let mut new_centroids = vec![0.0f32; self.config.n_clusters * dimension];
            cluster_sizes = vec![0usize; self.config.n_clusters];

            for (i, &cluster) in assignments.iter().enumerate() {
                let c = cluster as usize;
                cluster_sizes[c] += 1;
                let v_start = i * dimension;
                let c_start = c * dimension;
                for d in 0..dimension {
                    new_centroids[c_start + d] += vectors[v_start + d];
                }
            }

            // Divide by cluster size (handle empty clusters)
            for c in 0..self.config.n_clusters {
                if cluster_sizes[c] > 0 {
                    let c_start = c * dimension;
                    let size = cluster_sizes[c] as f32;
                    for d in 0..dimension {
                        new_centroids[c_start + d] /= size;
                    }
                } else {
                    // Reinitialize empty cluster with a random vector
                    let random_idx = (iter * 7 + c * 13) % n_vectors;
                    let v_start = random_idx * dimension;
                    let c_start = c * dimension;
                    new_centroids[c_start..c_start + dimension]
                        .copy_from_slice(&vectors[v_start..v_start + dimension]);
                }
            }

            // Transfer new centroids back to GPU
            centroids = Tensor::from_slice(
                &new_centroids,
                (self.config.n_clusters, dimension),
                &device,
            )
            .map_err(|e| GpuError::TransferError(format!("upload centroids: {e}")))?;
        }

        // Release memory pool reservation
        if let Some(pool) = memory_pool {
            pool.release(working_memory);
        }

        // Final centroid extraction
        let final_centroids = centroids
            .to_device(&Device::Cpu)
            .map_err(|e| GpuError::TransferError(format!("copy final centroids: {e}")))?
            .flatten_all()
            .map_err(|e| GpuError::TransferError(format!("flatten centroids: {e}")))?
            .to_vec1::<f32>()
            .map_err(|e| GpuError::TransferError(format!("convert centroids: {e}")))?;

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

        tracing::info!(
            "GPU KMeans completed: {} clusters, {} iterations, {}{}, {:.1}ms",
            self.config.n_clusters,
            iterations,
            if converged { "converged" } else { "max_iter" },
            if converged { "" } else { " (not converged)" },
            elapsed
        );

        Ok(KMeansResult {
            centroids: final_centroids,
            n_clusters: self.config.n_clusters,
            dimension,
            assignments,
            cluster_sizes,
            inertia: prev_inertia,
            iterations,
            converged,
            used_gpu: true,
            train_time_ms: elapsed,
        })
    }

    /// KMeans++ initialization on GPU.
    #[cfg(feature = "candle")]
    fn kmeans_plus_plus_init_gpu(
        &self,
        vectors: &candle_core::Tensor,
        n_vectors: usize,
        dimension: usize,
        device: &candle_core::Device,
    ) -> GpuResult<candle_core::Tensor> {
        use candle_core::{Tensor, DType};

        // Pick first centroid randomly (based on seed)
        let first_idx = (self.config.seed as usize) % n_vectors;
        let first = vectors
            .narrow(0, first_idx, 1)
            .map_err(|e| GpuError::KernelError(format!("narrow: {e}")))?;

        let mut centroid_list: Vec<Vec<f32>> = vec![
            first
                .to_device(&candle_core::Device::Cpu)
                .map_err(|e| GpuError::TransferError(format!("copy first centroid: {e}")))?
                .flatten_all()
                .map_err(|e| GpuError::KernelError(format!("flatten: {e}")))?
                .to_vec1::<f32>()
                .map_err(|e| GpuError::TransferError(format!("convert: {e}")))?
        ];

        for k in 1..self.config.n_clusters {
            // Build centroids tensor from selected so far
            let flat: Vec<f32> = centroid_list.iter().flatten().copied().collect();
            let current_centroids = Tensor::from_slice(&flat, (k, dimension), device)
                .map_err(|e| GpuError::TransferError(format!("upload centroids: {e}")))?;

            // Compute distances from all vectors to nearest existing centroid
            // Using ||a-b||² = ||a||² + ||b||² - 2aᵀb
            let dot = vectors
                .matmul(
                    &current_centroids
                        .t()
                        .map_err(|e| GpuError::KernelError(format!("transpose: {e}")))?,
                )
                .map_err(|e| GpuError::KernelError(format!("matmul: {e}")))?;

            let v_norms = vectors
                .sqr().map_err(|e| GpuError::KernelError(format!("sqr: {e}")))?
                .sum(1).map_err(|e| GpuError::KernelError(format!("sum: {e}")))?
                .unsqueeze(1).map_err(|e| GpuError::KernelError(format!("unsqueeze: {e}")))?;

            let c_norms = current_centroids
                .sqr().map_err(|e| GpuError::KernelError(format!("sqr: {e}")))?
                .sum(1).map_err(|e| GpuError::KernelError(format!("sum: {e}")))?
                .unsqueeze(0).map_err(|e| GpuError::KernelError(format!("unsqueeze: {e}")))?;

            let distances = dot
                .affine(-2.0, 0.0)
                .map_err(|e| GpuError::KernelError(format!("affine: {e}")))?
                .broadcast_add(&v_norms)
                .map_err(|e| GpuError::KernelError(format!("add: {e}")))?
                .broadcast_add(&c_norms)
                .map_err(|e| GpuError::KernelError(format!("add: {e}")))?;

            // Min distance to any centroid for each vector
            let min_dists = distances
                .min(1)
                .map_err(|e| GpuError::KernelError(format!("min: {e}")))?;

            // Transfer min distances to CPU for weighted random selection
            let min_dists_cpu: Vec<f32> = min_dists
                .to_device(&candle_core::Device::Cpu)
                .map_err(|e| GpuError::TransferError(format!("copy dists: {e}")))?
                .to_vec1::<f32>()
                .map_err(|e| GpuError::TransferError(format!("convert dists: {e}")))?;

            // Weighted selection (proportional to distance²)
            let total: f64 = min_dists_cpu.iter().map(|d| (*d).max(0.0) as f64).sum();
            if total <= 0.0 {
                // All distances are 0 — degenerate case, pick random
                let idx = (self.config.seed as usize + k * 997) % n_vectors;
                let row = vectors
                    .narrow(0, idx, 1)
                    .map_err(|e| GpuError::KernelError(format!("narrow: {e}")))?;
                centroid_list.push(
                    row.to_device(&candle_core::Device::Cpu)
                        .map_err(|e| GpuError::TransferError(format!("copy: {e}")))?
                        .flatten_all()
                        .map_err(|e| GpuError::KernelError(format!("flatten: {e}")))?
                        .to_vec1::<f32>()
                        .map_err(|e| GpuError::TransferError(format!("convert: {e}")))?,
                );
                continue;
            }

            // Deterministic weighted selection (no rand dependency, use seed)
            let threshold = ((self.config.seed as f64 + k as f64 * 3.14159) % 1.0) * total;
            let mut cumsum = 0.0;
            let mut selected_idx = 0;
            for (i, d) in min_dists_cpu.iter().enumerate() {
                cumsum += (*d).max(0.0) as f64;
                if cumsum >= threshold {
                    selected_idx = i;
                    break;
                }
            }

            let row = vectors
                .narrow(0, selected_idx, 1)
                .map_err(|e| GpuError::KernelError(format!("narrow: {e}")))?;
            centroid_list.push(
                row.to_device(&candle_core::Device::Cpu)
                    .map_err(|e| GpuError::TransferError(format!("copy: {e}")))?
                    .flatten_all()
                    .map_err(|e| GpuError::KernelError(format!("flatten: {e}")))?
                    .to_vec1::<f32>()
                    .map_err(|e| GpuError::TransferError(format!("convert: {e}")))?,
            );
        }

        // Assemble final centroids tensor on GPU
        let flat: Vec<f32> = centroid_list.into_iter().flatten().collect();
        Tensor::from_slice(&flat, (self.config.n_clusters, dimension), device)
            .map_err(|e| GpuError::TransferError(format!("upload final centroids: {e}")))
    }

    /// Random initialization on GPU.
    #[cfg(feature = "candle")]
    fn random_init_gpu(
        &self,
        vectors: &candle_core::Tensor,
        dimension: usize,
        device: &candle_core::Device,
    ) -> GpuResult<candle_core::Tensor> {
        let n_vectors = vectors.dim(0)
            .map_err(|e| GpuError::KernelError(format!("dim: {e}")))?;

        // Select K random indices (deterministic based on seed)
        let mut indices = Vec::with_capacity(self.config.n_clusters);
        let mut idx = self.config.seed as usize;
        for _ in 0..self.config.n_clusters {
            idx = (idx * 6364136223846793005 + 1442695040888963407) % n_vectors;
            indices.push(idx);
        }

        // Gather centroids from vectors
        let mut centroid_data = Vec::with_capacity(self.config.n_clusters * dimension);
        let vectors_cpu = vectors
            .to_device(&candle_core::Device::Cpu)
            .map_err(|e| GpuError::TransferError(format!("copy vectors: {e}")))?;

        for &idx in &indices {
            let row = vectors_cpu
                .narrow(0, idx, 1)
                .map_err(|e| GpuError::KernelError(format!("narrow: {e}")))?
                .flatten_all()
                .map_err(|e| GpuError::KernelError(format!("flatten: {e}")))?
                .to_vec1::<f32>()
                .map_err(|e| GpuError::TransferError(format!("convert: {e}")))?;
            centroid_data.extend_from_slice(&row);
        }

        candle_core::Tensor::from_slice(
            &centroid_data,
            (self.config.n_clusters, dimension),
            device,
        )
        .map_err(|e| GpuError::TransferError(format!("upload centroids: {e}")))
    }

    /// CPU fallback KMeans training (rayon-parallelized).
    fn train_cpu(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
    ) -> GpuResult<KMeansResult> {
        use rayon::prelude::*;

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

        // KMeans++ initialization on CPU
        let mut centroids = self.kmeans_plus_plus_init_cpu(vectors, n_vectors, dimension);

        let mut assignments = vec![0u32; n_vectors];
        let mut cluster_sizes = vec![0usize; self.config.n_clusters];
        let mut prev_inertia = f64::MAX;
        let mut iterations = 0;
        let mut converged = false;

        for iter in 0..self.config.max_iterations {
            iterations = iter + 1;

            // Assignment step (parallel)
            let results: Vec<(u32, f32)> = (0..n_vectors)
                .into_par_iter()
                .map(|i| {
                    let v = &vectors[i * dimension..(i + 1) * dimension];
                    let mut best_cluster = 0u32;
                    let mut best_dist = f32::MAX;

                    for c in 0..self.config.n_clusters {
                        let cent = &centroids[c * dimension..(c + 1) * dimension];
                        let dist: f32 = v.iter().zip(cent).map(|(a, b)| (a - b) * (a - b)).sum();
                        if dist < best_dist {
                            best_dist = dist;
                            best_cluster = c as u32;
                        }
                    }
                    (best_cluster, best_dist)
                })
                .collect();

            let inertia: f64 = results.iter().map(|(_, d)| *d as f64).sum();
            for (i, (cluster, _)) in results.iter().enumerate() {
                assignments[i] = *cluster;
            }

            // Convergence check
            let inertia_change = ((prev_inertia - inertia) / prev_inertia).abs();
            if self.config.verbose {
                tracing::info!(
                    "CPU KMeans iter {}/{}: inertia={:.2}, change={:.6}",
                    iterations, self.config.max_iterations, inertia, inertia_change
                );
            }

            if inertia_change < self.config.tolerance && iter > 0 {
                converged = true;
                prev_inertia = inertia;
                break;
            }
            prev_inertia = inertia;

            // Update step
            centroids = vec![0.0f32; self.config.n_clusters * dimension];
            cluster_sizes = vec![0usize; self.config.n_clusters];

            for (i, &cluster) in assignments.iter().enumerate() {
                let c = cluster as usize;
                cluster_sizes[c] += 1;
                let v_start = i * dimension;
                let c_start = c * dimension;
                for d in 0..dimension {
                    centroids[c_start + d] += vectors[v_start + d];
                }
            }

            for c in 0..self.config.n_clusters {
                if cluster_sizes[c] > 0 {
                    let c_start = c * dimension;
                    let size = cluster_sizes[c] as f32;
                    for d in 0..dimension {
                        centroids[c_start + d] /= size;
                    }
                } else {
                    let random_idx = (iter * 7 + c * 13) % n_vectors;
                    let v_start = random_idx * dimension;
                    let c_start = c * dimension;
                    centroids[c_start..c_start + dimension]
                        .copy_from_slice(&vectors[v_start..v_start + dimension]);
                }
            }
        }

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

        Ok(KMeansResult {
            centroids,
            n_clusters: self.config.n_clusters,
            dimension,
            assignments,
            cluster_sizes,
            inertia: prev_inertia,
            iterations,
            converged,
            used_gpu: false,
            train_time_ms: elapsed,
        })
    }

    /// KMeans++ initialization on CPU.
    fn kmeans_plus_plus_init_cpu(
        &self,
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
    ) -> Vec<f32> {
        let mut centroids = Vec::with_capacity(self.config.n_clusters * dimension);

        // First centroid: based on seed
        let first_idx = (self.config.seed as usize) % n_vectors;
        centroids.extend_from_slice(&vectors[first_idx * dimension..(first_idx + 1) * dimension]);

        let mut min_distances = vec![f32::MAX; n_vectors];

        for k in 1..self.config.n_clusters {
            // Update min distances with last added centroid
            let last_centroid = &centroids[(k - 1) * dimension..k * dimension];
            for i in 0..n_vectors {
                let v = &vectors[i * dimension..(i + 1) * dimension];
                let dist: f32 = v.iter().zip(last_centroid).map(|(a, b)| (a - b) * (a - b)).sum();
                min_distances[i] = min_distances[i].min(dist);
            }

            // Weighted selection
            let total: f64 = min_distances.iter().map(|d| *d as f64).sum();
            if total <= 0.0 {
                let idx = (self.config.seed as usize + k * 997) % n_vectors;
                centroids.extend_from_slice(&vectors[idx * dimension..(idx + 1) * dimension]);
                continue;
            }

            let threshold = ((self.config.seed as f64 + k as f64 * 3.14159) % 1.0) * total;
            let mut cumsum = 0.0;
            let mut selected = 0;
            for (i, d) in min_distances.iter().enumerate() {
                cumsum += *d as f64;
                if cumsum >= threshold {
                    selected = i;
                    break;
                }
            }

            centroids.extend_from_slice(&vectors[selected * dimension..(selected + 1) * dimension]);
        }

        centroids
    }
}

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

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

    fn generate_clustered_data(
        n_clusters: usize,
        points_per_cluster: usize,
        dimension: usize,
    ) -> Vec<f32> {
        let mut data = Vec::new();
        for c in 0..n_clusters {
            let center = c as f32 * 10.0; // Clusters well-separated
            for p in 0..points_per_cluster {
                for d in 0..dimension {
                    // Small noise around cluster center
                    let noise = ((p * 7 + d * 13 + c * 31) % 100) as f32 / 100.0 - 0.5;
                    data.push(center + noise);
                }
            }
        }
        data
    }

    #[test]
    fn test_cpu_kmeans_basic() {
        let n_clusters = 3;
        let points_per_cluster = 100;
        let dimension = 4;
        let data = generate_clustered_data(n_clusters, points_per_cluster, dimension);
        let n_vectors = n_clusters * points_per_cluster;

        let config = GpuKMeansConfig {
            n_clusters,
            max_iterations: 50,
            tolerance: 1e-4,
            seed: 42,
            verbose: false,
            ..Default::default()
        };

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_cpu(&data, n_vectors, dimension).unwrap();

        assert_eq!(result.n_clusters, n_clusters);
        assert_eq!(result.dimension, dimension);
        assert_eq!(result.assignments.len(), n_vectors);
        assert_eq!(result.centroids.len(), n_clusters * dimension);
        assert!(!result.used_gpu);
        assert!(result.inertia >= 0.0);
        assert!(result.iterations > 0);

        // Each cluster should have approximately points_per_cluster members
        for &size in &result.cluster_sizes {
            assert!(size > 0, "No empty clusters expected");
        }
    }

    #[test]
    fn test_cpu_kmeans_convergence() {
        let data = generate_clustered_data(2, 200, 8);
        let config = GpuKMeansConfig {
            n_clusters: 2,
            max_iterations: 100,
            tolerance: 1e-6,
            seed: 123,
            ..Default::default()
        };

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_cpu(&data, 400, 8).unwrap();

        // Should converge with well-separated clusters
        assert!(result.converged || result.iterations <= 100);
    }

    #[test]
    fn test_kmeans_plus_plus_init() {
        let data = generate_clustered_data(4, 50, 2);
        let config = GpuKMeansConfig {
            n_clusters: 4,
            seed: 42,
            ..Default::default()
        };

        let trainer = GpuKMeansTrainer::new(config);
        let centroids = trainer.kmeans_plus_plus_init_cpu(&data, 200, 2);

        assert_eq!(centroids.len(), 4 * 2);

        // Centroids should be distinct
        for i in 0..4 {
            for j in (i + 1)..4 {
                let ci = &centroids[i * 2..(i + 1) * 2];
                let cj = &centroids[j * 2..(j + 1) * 2];
                let dist: f32 = ci.iter().zip(cj).map(|(a, b)| (a - b) * (a - b)).sum();
                assert!(dist > 0.0, "Centroids {i} and {j} should be distinct");
            }
        }
    }

    #[test]
    fn test_kmeans_result_centroid_access() {
        let result = KMeansResult {
            centroids: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            n_clusters: 2,
            dimension: 3,
            assignments: vec![0, 1],
            cluster_sizes: vec![1, 1],
            inertia: 0.0,
            iterations: 1,
            converged: true,
            used_gpu: false,
            train_time_ms: 0,
        };

        assert_eq!(result.centroid(0), Some(&[1.0, 2.0, 3.0][..]));
        assert_eq!(result.centroid(1), Some(&[4.0, 5.0, 6.0][..]));
        assert_eq!(result.centroid(2), None);
    }

    #[test]
    fn test_kmeans_input_validation() {
        let trainer = GpuKMeansTrainer::new(GpuKMeansConfig {
            n_clusters: 10,
            ..Default::default()
        });

        // Empty input
        let result = trainer.train(&[], 0, 0, None);
        assert!(result.is_err());

        // More clusters than vectors
        let data = vec![0.0f32; 3 * 4]; // 3 vectors of dim 4
        let result = trainer.train(&data, 3, 4, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_kmeans_train_with_fallback() {
        // This test exercises the full train() path which tries GPU then falls back to CPU
        let data = generate_clustered_data(3, 100, 4);
        let config = GpuKMeansConfig {
            n_clusters: 3,
            max_iterations: 20,
            seed: 42,
            ..Default::default()
        };

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train(&data, 300, 4, None).unwrap();

        assert_eq!(result.n_clusters, 3);
        assert_eq!(result.assignments.len(), 300);
        // On CI/machines without GPU, should fall back to CPU
        // (used_gpu may be true or false depending on environment)
    }

    // =====================================================================
    // Phase-63: DeviceSelector integration tests
    // =====================================================================

    #[test]
    fn test_kmeans_with_cpu_selector() {
        // Explicitly force CPU via DeviceSelector
        let data = generate_clustered_data(3, 100, 4);
        let config = GpuKMeansConfig {
            n_clusters: 3,
            max_iterations: 20,
            seed: 42,
            ..Default::default()
        };

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

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_with_selector(&data, 300, 4, None, Some(&selector)).unwrap();

        assert_eq!(result.n_clusters, 3);
        assert_eq!(result.assignments.len(), 300);
        assert!(!result.used_gpu, "CPU selector should force CPU path");
    }

    #[test]
    fn test_kmeans_with_auto_selector() {
        // Auto selector — will use GPU if available, otherwise CPU
        let data = generate_clustered_data(4, 500, 8);
        let config = GpuKMeansConfig {
            n_clusters: 4,
            max_iterations: 30,
            seed: 42,
            ..Default::default()
        };

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

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_with_selector(&data, 2000, 8, None, Some(&selector)).unwrap();

        assert_eq!(result.n_clusters, 4);
        assert_eq!(result.assignments.len(), 2000);
        assert!(result.centroids.len() == 4 * 8);
    }

    #[test]
    fn test_kmeans_small_data_skips_gpu() {
        // Very small dataset — should skip GPU even with Auto selector
        let data = generate_clustered_data(2, 10, 4);
        let config = GpuKMeansConfig {
            n_clusters: 2,
            max_iterations: 10,
            seed: 42,
            ..Default::default()
        };

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

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_with_selector(&data, 20, 4, None, Some(&selector)).unwrap();

        assert_eq!(result.n_clusters, 2);
        // Small data should use CPU even with auto selector
        assert!(!result.used_gpu, "Small data should use CPU path");
    }

    #[test]
    fn test_kmeans_benchmark_cpu() {
        // CPU performance baseline benchmark (not a real perf test, just ensures it runs)
        let n_clusters = 8;
        let n_vectors = 5000;
        let dimension = 32;
        let data = generate_clustered_data(n_clusters, n_vectors / n_clusters, dimension);

        let config = GpuKMeansConfig {
            n_clusters,
            max_iterations: 20,
            seed: 42,
            verbose: false,
            ..Default::default()
        };

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

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_with_selector(&data, n_vectors, dimension, None, Some(&selector)).unwrap();

        assert_eq!(result.n_clusters, n_clusters);
        assert!(!result.used_gpu);
        println!(
            "CPU KMeans benchmark: {} vectors × {}d, {} clusters, {} iters, {}ms",
            n_vectors, dimension, n_clusters, result.iterations, result.train_time_ms,
        );

        // Verify result quality — each cluster should have at least some members
        for (i, &size) in result.cluster_sizes.iter().enumerate() {
            assert!(size > 0, "Cluster {i} should not be empty");
        }
    }

    #[test]
    fn test_kmeans_auto_benchmark() {
        // Auto benchmark — GPU or CPU depending on environment
        let n_clusters = 4;
        let n_vectors = 2000;
        let dimension = 16;
        let data = generate_clustered_data(n_clusters, n_vectors / n_clusters, dimension);

        let config = GpuKMeansConfig {
            n_clusters,
            max_iterations: 30,
            seed: 42,
            verbose: false,
            ..Default::default()
        };

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

        let trainer = GpuKMeansTrainer::new(config);
        let result = trainer.train_with_selector(&data, n_vectors, dimension, None, Some(&selector)).unwrap();

        println!(
            "Auto KMeans benchmark: {} vectors × {}d, {} clusters, {} iters, {}ms, gpu={}",
            n_vectors, dimension, n_clusters, result.iterations, result.train_time_ms, result.used_gpu,
        );

        assert_eq!(result.n_clusters, n_clusters);
        assert!(result.centroids.len() == n_clusters * dimension);
    }
}
