//! CAGRA graph index — GPU-exclusive approximate nearest neighbor graph
//!
//! CAGRA (Compressed Adaptive Graph for Rapid Approximation) is a GPU-native
//! graph index algorithm proposed by NVIDIA in 2023. It is designed to be
//! 10-100x faster than HNSW for both build and search on GPU hardware.
//!
//! ## Algorithm Overview
//!
//! 1. **Build phase** (GPU):
//!    - Construct a k-NN graph using brute-force on GPU (fast due to parallelism)
//!    - Prune the graph to `graph_degree` edges per node
//!    - Optimize graph connectivity with reverse-edge addition
//!
//! 2. **Search phase** (GPU):
//!    - Multi-start greedy traversal on the graph
//!    - Each search thread explores different entry points
//!    - Merge results from all threads
//!
//! ## CPU Fallback
//!
//! When GPU is not available, CAGRA falls back to building an HNSW-style
//! graph on CPU. The graph structure is compatible for search on either device.

use super::error::{GpuError, GpuResult};
use super::memory::GpuMemoryPool;
use super::batch_builder::{StreamingBatchBuilder, BatchBuilderConfig};

/// CAGRA graph index configuration.
#[derive(Debug, Clone)]
pub struct CagraConfig {
    /// Graph degree: number of edges per node in the final graph.
    /// Higher values → better recall but larger index. Default: 64.
    pub graph_degree: usize,

    /// Intermediate graph degree during build (before pruning).
    /// Should be >= graph_degree. Default: 128.
    pub intermediate_graph_degree: usize,

    /// Number of iterations for graph optimization. Default: 10.
    pub build_iterations: usize,

    /// Whether to add reverse edges for better connectivity.
    pub add_reverse_edges: bool,

    /// Search width: number of candidates explored in parallel during search.
    /// Higher → better recall but slower. Default: 32.
    pub search_width: usize,

    /// Min/max iterations for search termination. Default: 0 (auto).
    pub min_iterations: usize,
    pub max_iterations: usize,

    /// Enable verbose logging.
    pub verbose: bool,
}

impl Default for CagraConfig {
    fn default() -> Self {
        Self {
            graph_degree: 64,
            intermediate_graph_degree: 128,
            build_iterations: 10,
            add_reverse_edges: true,
            search_width: 32,
            min_iterations: 0,
            max_iterations: 0,
            verbose: false,
        }
    }
}

/// CAGRA graph index — an in-memory graph for GPU-accelerated ANN search.
///
/// The graph stores edges for each node: `graph[node_id]` → list of neighbor IDs.
/// Vectors are stored separately (in the original Lance dataset).
pub struct CagraIndex {
    /// Adjacency list: graph[i] = neighbor IDs for node i (sorted by distance).
    /// Each node has exactly `graph_degree` neighbors.
    graph: Vec<Vec<u32>>,

    /// Number of vectors in the index
    num_vectors: usize,

    /// Vector dimension
    dimension: usize,

    /// Graph degree (edges per node)
    graph_degree: usize,

    /// Whether the graph was built on GPU
    built_on_gpu: bool,

    /// Build time in milliseconds
    build_time_ms: u64,
}

/// Result of a CAGRA search.
#[derive(Debug, Clone)]
pub struct CagraSearchResult {
    /// Neighbor indices (sorted by distance, nearest first)
    pub indices: Vec<u32>,
    /// Distances to query (same order as indices)
    pub distances: Vec<f32>,
}

/// Build statistics for CAGRA index.
#[derive(Debug, Clone)]
pub struct CagraBuildStats {
    /// Number of vectors indexed
    pub num_vectors: usize,
    /// Vector dimension
    pub dimension: usize,
    /// Graph degree
    pub graph_degree: usize,
    /// Build time in milliseconds
    pub build_time_ms: u64,
    /// Whether GPU was used
    pub used_gpu: bool,
    /// Peak memory usage in bytes (approximate)
    pub peak_memory_bytes: usize,
    /// Average out-degree (should be close to graph_degree)
    pub avg_out_degree: f64,
}

impl CagraIndex {
    /// Build a CAGRA graph index from vectors.
    ///
    /// Tries GPU-accelerated build first, falls back to CPU if unavailable.
    pub fn build(
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        config: &CagraConfig,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> GpuResult<Self> {
        if n_vectors == 0 {
            return Err(GpuError::Other("Empty vectors for CAGRA build".into()));
        }
        if vectors.len() != n_vectors * dimension {
            return Err(GpuError::Other("Vector data size mismatch".into()));
        }
        if n_vectors < config.graph_degree {
            return Err(GpuError::Other(format!(
                "Too few vectors ({}) for graph_degree ({})",
                n_vectors, config.graph_degree
            )));
        }

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

        // Try GPU build
        #[cfg(feature = "candle")]
        {
            match Self::build_gpu(vectors, n_vectors, dimension, config, memory_pool) {
                Ok(index) => return Ok(index),
                Err(e) => {
                    tracing::warn!("GPU CAGRA build failed, falling back to CPU: {}", e);
                }
            }
        }

        // CPU fallback: build a simple k-NN graph
        let graph = Self::build_knn_graph_cpu(vectors, n_vectors, dimension, config);
        let elapsed = start.elapsed().as_millis() as u64;

        tracing::info!(
            "CAGRA CPU build: {} vectors, {}d, degree={}, {:.1}ms",
            n_vectors, dimension, config.graph_degree, elapsed
        );

        Ok(Self {
            graph,
            num_vectors: n_vectors,
            dimension,
            graph_degree: config.graph_degree,
            built_on_gpu: false,
            build_time_ms: elapsed,
        })
    }

    /// GPU-accelerated CAGRA graph build using brute-force k-NN.
    #[cfg(feature = "candle")]
    fn build_gpu(
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        config: &CagraConfig,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> GpuResult<Self> {
        use candle_core::{Device, Tensor};

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

        let device = Device::new_cuda(0).map_err(|e| {
            GpuError::InitializationError(format!("CUDA init failed: {e}"))
        })?;

        // Check memory: need vectors + distance matrix (for batches)
        let data_bytes = n_vectors * dimension * 4;
        if let Some(pool) = memory_pool {
            if !pool.can_fit(data_bytes * 2) {
                return Err(GpuError::OutOfMemory {
                    requested_bytes: data_bytes * 2,
                    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!("upload vectors: {e}")))?;

        // Build k-NN graph using batched brute-force
        let k = config.intermediate_graph_degree.min(n_vectors - 1);
        let mut graph = vec![Vec::with_capacity(config.graph_degree); n_vectors];

        // Process in batches to avoid OOM on the distance matrix
        let batch_size = Self::compute_knn_batch_size(n_vectors, dimension, memory_pool);
        let num_batches = (n_vectors + batch_size - 1) / batch_size;

        for b in 0..num_batches {
            let batch_start = b * batch_size;
            let batch_end = (batch_start + batch_size).min(n_vectors);
            let batch_count = batch_end - batch_start;

            // Extract batch from GPU tensor
            let batch_tensor = vectors_tensor.narrow(0, batch_start, batch_count)
                .map_err(|e| GpuError::KernelError(format!("narrow batch: {e}")))?;

            // Compute all-pairs distances: batch × all vectors
            // Using: ||a-b||² = ||a||² + ||b||² - 2aᵀb
            let dot = batch_tensor
                .matmul(&vectors_tensor.t()
                    .map_err(|e| GpuError::KernelError(format!("transpose: {e}")))?)
                .map_err(|e| GpuError::KernelError(format!("matmul: {e}")))?;

            let batch_norms = batch_tensor
                .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 all_norms = vectors_tensor
                .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(&batch_norms)
                .map_err(|e| GpuError::KernelError(format!("add: {e}")))?
                .broadcast_add(&all_norms)
                .map_err(|e| GpuError::KernelError(format!("add: {e}")))?;

            // Transfer distances to CPU for top-k selection
            let dist_cpu: Vec<f32> = distances
                .to_device(&Device::Cpu)
                .map_err(|e| GpuError::TransferError(format!("copy distances: {e}")))?
                .flatten_all()
                .map_err(|e| GpuError::KernelError(format!("flatten: {e}")))?
                .to_vec1::<f32>()
                .map_err(|e| GpuError::TransferError(format!("convert distances: {e}")))?;

            // For each query in batch, find top-k nearest neighbors
            for qi in 0..batch_count {
                let global_idx = batch_start + qi;
                let row_start = qi * n_vectors;

                // Collect (distance, index) pairs, excluding self
                let mut candidates: Vec<(f32, u32)> = (0..n_vectors)
                    .filter(|&j| j != global_idx)
                    .map(|j| (dist_cpu[row_start + j].max(0.0), j as u32))
                    .collect();

                // Partial sort to get top-k
                candidates.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
                let top_k: Vec<u32> = candidates.iter()
                    .take(k)
                    .map(|&(_, idx)| idx)
                    .collect();

                graph[global_idx] = top_k;
            }

            if config.verbose && (b % 10 == 0 || b == num_batches - 1) {
                tracing::info!(
                    "CAGRA GPU k-NN: batch {}/{}, {:.1}%",
                    b + 1, num_batches,
                    (batch_end as f64 / n_vectors as f64) * 100.0
                );
            }
        }

        // Prune to graph_degree
        if config.graph_degree < k {
            for neighbors in &mut graph {
                neighbors.truncate(config.graph_degree);
            }
        }

        // Add reverse edges for better connectivity
        if config.add_reverse_edges {
            Self::add_reverse_edges(&mut graph, config.graph_degree);
        }

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

        tracing::info!(
            "CAGRA GPU build: {} vectors, {}d, degree={}, {:.1}ms",
            n_vectors, dimension, config.graph_degree, elapsed
        );

        Ok(Self {
            graph,
            num_vectors: n_vectors,
            dimension,
            graph_degree: config.graph_degree,
            built_on_gpu: true,
            build_time_ms: elapsed,
        })
    }

    /// Compute optimal batch size for k-NN computation.
    fn compute_knn_batch_size(
        n_vectors: usize,
        dimension: usize,
        memory_pool: Option<&GpuMemoryPool>,
    ) -> usize {
        // Distance matrix for batch: batch_size × n_vectors × 4 bytes
        let available = memory_pool
            .map(|p| p.budget().saturating_sub(p.allocated()) / 4) // Leave room
            .unwrap_or(256 * 1024 * 1024); // 256MB default

        let row_size = n_vectors * std::mem::size_of::<f32>();
        if row_size == 0 {
            return n_vectors;
        }

        let max_batch = available / row_size;
        max_batch.max(64).min(n_vectors) // At least 64, at most all
    }

    /// Build k-NN graph on CPU (fallback).
    fn build_knn_graph_cpu(
        vectors: &[f32],
        n_vectors: usize,
        dimension: usize,
        config: &CagraConfig,
    ) -> Vec<Vec<u32>> {
        use rayon::prelude::*;

        let k = config.graph_degree;

        let graph: Vec<Vec<u32>> = (0..n_vectors)
            .into_par_iter()
            .map(|i| {
                let v = &vectors[i * dimension..(i + 1) * dimension];

                // Compute distances to all other vectors
                let mut candidates: Vec<(f32, u32)> = (0..n_vectors)
                    .filter(|&j| j != i)
                    .map(|j| {
                        let u = &vectors[j * dimension..(j + 1) * dimension];
                        let dist: f32 = v.iter().zip(u).map(|(a, b)| (a - b) * (a - b)).sum();
                        (dist, j as u32)
                    })
                    .collect();

                // Partial sort for top-k
                candidates.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
                candidates.into_iter().take(k).map(|(_, idx)| idx).collect()
            })
            .collect();

        graph
    }

    /// Add reverse edges to improve graph connectivity.
    ///
    /// For each edge A→B, also adds B→A (if B doesn't already have max neighbors).
    fn add_reverse_edges(graph: &mut [Vec<u32>], max_degree: usize) {
        let n = graph.len();
        let mut reverse_edges: Vec<Vec<u32>> = vec![Vec::new(); n];

        // Collect reverse edges
        for (i, neighbors) in graph.iter().enumerate() {
            for &j in neighbors {
                let j = j as usize;
                if j < n && !graph[j].contains(&(i as u32)) {
                    reverse_edges[j].push(i as u32);
                }
            }
        }

        // Add reverse edges (up to max_degree)
        for (i, rev_edges) in reverse_edges.into_iter().enumerate() {
            for edge in rev_edges {
                if graph[i].len() < max_degree && !graph[i].contains(&edge) {
                    graph[i].push(edge);
                }
            }
        }
    }

    /// Search the CAGRA graph for nearest neighbors.
    ///
    /// Uses greedy graph traversal starting from random entry points.
    pub fn search(
        &self,
        query: &[f32],
        k: usize,
        vectors: &[f32],
        search_width: Option<usize>,
    ) -> CagraSearchResult {
        let width = search_width.unwrap_or(32).max(k);
        let dimension = self.dimension;

        // Multi-start greedy search
        let num_starts = width.min(self.num_vectors);

        // Use deterministic entry points (evenly spaced)
        let step = if num_starts > 0 { self.num_vectors / num_starts } else { 1 };
        let entry_points: Vec<usize> = (0..num_starts).map(|i| i * step).collect();

        // Best-first search
        let mut visited = vec![false; self.num_vectors];
        let mut candidates: Vec<(f32, u32)> = Vec::with_capacity(width * 2);

        // Initialize with entry points
        for &ep in &entry_points {
            if ep < self.num_vectors && !visited[ep] {
                visited[ep] = true;
                let dist = Self::l2_distance(
                    query,
                    &vectors[ep * dimension..(ep + 1) * dimension],
                );
                candidates.push((dist, ep as u32));
            }
        }

        candidates.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());

        // Greedy expansion
        let max_iters = self.num_vectors.min(width * 10);
        let mut expanded = 0;

        while expanded < candidates.len() && expanded < max_iters {
            let (_, node) = candidates[expanded];
            expanded += 1;

            // Visit neighbors
            if let Some(neighbors) = self.graph.get(node as usize) {
                for &neighbor in neighbors {
                    let n = neighbor as usize;
                    if n < self.num_vectors && !visited[n] {
                        visited[n] = true;
                        let dist = Self::l2_distance(
                            query,
                            &vectors[n * dimension..(n + 1) * dimension],
                        );
                        candidates.push((dist, neighbor));
                    }
                }
            }

            // Keep top candidates sorted
            if candidates.len() > width * 4 {
                candidates.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
                candidates.truncate(width * 2);
            }
        }

        // Final sort and return top-k
        candidates.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
        candidates.truncate(k);

        CagraSearchResult {
            indices: candidates.iter().map(|&(_, idx)| idx).collect(),
            distances: candidates.iter().map(|&(dist, _)| dist).collect(),
        }
    }

    /// L2 squared distance between two vectors.
    #[inline]
    fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
        a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum()
    }

    /// Get build statistics.
    pub fn build_stats(&self) -> CagraBuildStats {
        let total_edges: usize = self.graph.iter().map(|n| n.len()).sum();
        let avg_degree = if self.num_vectors > 0 {
            total_edges as f64 / self.num_vectors as f64
        } else {
            0.0
        };

        CagraBuildStats {
            num_vectors: self.num_vectors,
            dimension: self.dimension,
            graph_degree: self.graph_degree,
            build_time_ms: self.build_time_ms,
            used_gpu: self.built_on_gpu,
            peak_memory_bytes: self.num_vectors * self.graph_degree * 4, // Approximate
            avg_out_degree: avg_degree,
        }
    }

    /// Number of vectors in the index.
    pub fn len(&self) -> usize {
        self.num_vectors
    }

    /// Whether the index is empty.
    pub fn is_empty(&self) -> bool {
        self.num_vectors == 0
    }

    /// Whether the index was built on GPU.
    pub fn built_on_gpu(&self) -> bool {
        self.built_on_gpu
    }

    /// Get the graph degree.
    pub fn graph_degree(&self) -> usize {
        self.graph_degree
    }
}

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

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

    fn generate_clustered_data(n_clusters: usize, per_cluster: usize, dim: usize) -> Vec<f32> {
        let mut data = Vec::new();
        for c in 0..n_clusters {
            let center = c as f32 * 20.0;
            for p in 0..per_cluster {
                for d in 0..dim {
                    let noise = ((p * 7 + d * 13 + c * 31) % 100) as f32 / 100.0 - 0.5;
                    data.push(center + noise);
                }
            }
        }
        data
    }

    #[test]
    fn test_cagra_build_cpu() {
        let n_vectors = 200;
        let dim = 4;
        let data = generate_clustered_data(4, 50, dim);

        let config = CagraConfig {
            graph_degree: 8,
            intermediate_graph_degree: 16,
            ..Default::default()
        };

        let index = CagraIndex::build(&data, n_vectors, dim, &config, None).unwrap();

        assert_eq!(index.len(), n_vectors);
        assert_eq!(index.graph_degree(), 8);
        assert!(!index.is_empty());

        let stats = index.build_stats();
        assert_eq!(stats.num_vectors, n_vectors);
        assert!(stats.avg_out_degree > 0.0);
    }

    #[test]
    fn test_cagra_search() {
        let n_vectors = 100;
        let dim = 4;
        let data = generate_clustered_data(4, 25, dim);

        let config = CagraConfig {
            graph_degree: 8,
            ..Default::default()
        };

        let index = CagraIndex::build(&data, n_vectors, dim, &config, None).unwrap();

        // Search for nearest neighbors of the first vector
        let query = &data[0..dim];
        let result = index.search(query, 5, &data, None);

        assert_eq!(result.indices.len(), 5);
        assert_eq!(result.distances.len(), 5);

        // First result should be very close (same cluster)
        assert!(result.distances[0] < 2.0, "First result too far: {}", result.distances[0]);

        // Distances should be sorted
        for i in 1..result.distances.len() {
            assert!(
                result.distances[i] >= result.distances[i - 1],
                "Distances not sorted at {}: {} < {}",
                i, result.distances[i], result.distances[i - 1]
            );
        }
    }

    #[test]
    fn test_cagra_search_recall() {
        // Build index and verify it finds true nearest neighbors
        let n_vectors = 200;
        let dim = 4;
        let data = generate_clustered_data(4, 50, dim);

        let config = CagraConfig {
            graph_degree: 16,
            ..Default::default()
        };

        let index = CagraIndex::build(&data, n_vectors, dim, &config, None).unwrap();

        // Compute true nearest neighbor by brute force
        let query = &data[0..dim];
        let mut brute_force: Vec<(f32, usize)> = (1..n_vectors)
            .map(|i| {
                let v = &data[i * dim..(i + 1) * dim];
                let dist: f32 = query.iter().zip(v).map(|(a, b)| (a - b) * (a - b)).sum();
                (dist, i)
            })
            .collect();
        brute_force.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());

        let true_nn: Vec<u32> = brute_force.iter().take(10).map(|&(_, i)| i as u32).collect();

        // Search with graph
        let result = index.search(query, 10, &data, Some(64));

        // Count how many true NNs are in the graph result
        let hits: usize = result
            .indices
            .iter()
            .filter(|idx| true_nn.contains(idx))
            .count();

        // With graph_degree=16, we should get good recall for top-10
        // Allow some misses since this is approximate
        assert!(
            hits >= 5,
            "Recall too low: {}/{} (found {:?}, true {:?})",
            hits, 10, result.indices, true_nn
        );
    }

    #[test]
    fn test_cagra_reverse_edges() {
        let mut graph = vec![
            vec![1, 2],    // 0 → 1, 2
            vec![0, 2],    // 1 → 0, 2
            vec![],        // 2 → nothing
        ];

        CagraIndex::add_reverse_edges(&mut graph, 4);

        // Node 2 should now have reverse edges from 0 and 1
        assert!(graph[2].contains(&0));
        assert!(graph[2].contains(&1));
    }

    #[test]
    fn test_cagra_l2_distance() {
        let a = [1.0, 2.0, 3.0];
        let b = [4.0, 5.0, 6.0];
        let dist = CagraIndex::l2_distance(&a, &b);
        // (4-1)² + (5-2)² + (6-3)² = 9 + 9 + 9 = 27
        assert!((dist - 27.0).abs() < 1e-6);
    }

    #[test]
    fn test_cagra_build_validation() {
        let config = CagraConfig::default();

        // Empty vectors
        assert!(CagraIndex::build(&[], 0, 4, &config, None).is_err());

        // Too few vectors for graph_degree
        let data = vec![0.0f32; 4 * 2]; // Only 2 vectors
        let config = CagraConfig {
            graph_degree: 64,
            ..Default::default()
        };
        assert!(CagraIndex::build(&data, 2, 4, &config, None).is_err());
    }

    #[test]
    fn test_cagra_build_stats() {
        let data = generate_clustered_data(2, 50, 4);
        let config = CagraConfig {
            graph_degree: 4,
            ..Default::default()
        };

        let index = CagraIndex::build(&data, 100, 4, &config, None).unwrap();
        let stats = index.build_stats();

        assert_eq!(stats.num_vectors, 100);
        assert_eq!(stats.dimension, 4);
        assert_eq!(stats.graph_degree, 4);
        assert!(!stats.used_gpu); // No GPU in test environment
        assert!(stats.avg_out_degree >= 4.0); // At least graph_degree (reverse edges may add more)
    }
}
