lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! Hierarchical Navigable Small World (HNSW) graph for approximate nearest neighbor search.
//!
//! HNSW is a graph-based index structure that provides logarithmic search complexity
//! with high recall rates. It builds a multi-layer graph where:
//!
//! - Each layer is a navigable small world graph
//! - Higher layers contain fewer nodes (exponentially decreasing)
//! - Search starts from the top layer and greedily descends
//! - The bottom layer (layer 0) contains all nodes
//!
//! # Algorithm
//!
//! ## Insertion
//!
//! 1. Randomly select the maximum layer for the new node (exponential distribution)
//! 2. Start from the entry point at the top layer
//! 3. Greedily search for the nearest neighbor at each layer (layers > insertion layer)
//! 4. At and below the insertion layer, find M nearest neighbors and connect
//! 5. Update connections to maintain graph properties
//!
//! ## Search
//!
//! 1. Start from the entry point at the top layer
//! 2. At each layer, perform greedy search keeping `ef` candidates
//! 3. Move to the next layer using the best candidates as entry points
//! 4. At the bottom layer, return top-k from the final candidate set
//!
//! # Parameters
//!
//! - `M`: Maximum neighbors per node at layer 0 (default: 16)
//! - `M_max0`: Maximum neighbors at layer 0 (typically 2*M)
//! - `ef_construction`: Beam width during construction (default: 200)
//! - `ef_search`: Beam width during search (default: 50)
//! - `ml`: Level multiplier, controls layer probability (default: 1/ln(M))
//!
//! # References
//!
//! - Malkov, Y. A., & Yashunin, D. A. (2018). Efficient and robust approximate
//!   nearest neighbor search using Hierarchical Navigable Small World graphs.

use alloc::collections::BinaryHeap;
use alloc::vec::Vec;
use core::cmp::Ordering;
use hashbrown::HashMap;

use super::distance::{compute_distance, cosine_distance};
use super::types::{DistanceMetric, HNSW_MAX_NEIGHBORS, HnswNode, VectorError, VectorSearchResult};

// ═══════════════════════════════════════════════════════════════════════════════
// PRIORITY QUEUE HELPERS
// ═══════════════════════════════════════════════════════════════════════════════

/// Candidate node during search (max-heap by distance for efficient eviction).
#[derive(Clone)]
struct MaxCandidate {
    id: u64,
    distance: f32,
}

impl PartialEq for MaxCandidate {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for MaxCandidate {}

impl PartialOrd for MaxCandidate {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MaxCandidate {
    fn cmp(&self, other: &Self) -> Ordering {
        // Max-heap: larger distance = higher priority
        self.distance
            .partial_cmp(&other.distance)
            .unwrap_or(Ordering::Equal)
    }
}

/// Candidate node during search (min-heap by distance for nearest-first).
#[derive(Clone)]
struct MinCandidate {
    id: u64,
    distance: f32,
}

impl PartialEq for MinCandidate {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for MinCandidate {}

impl PartialOrd for MinCandidate {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MinCandidate {
    fn cmp(&self, other: &Self) -> Ordering {
        // Min-heap: smaller distance = higher priority (reverse comparison)
        other
            .distance
            .partial_cmp(&self.distance)
            .unwrap_or(Ordering::Equal)
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// HNSW LAYER
// ═══════════════════════════════════════════════════════════════════════════════

/// A single layer in the HNSW graph.
#[derive(Debug, Clone, Default)]
struct HnswLayer {
    /// Adjacency list: node_id -> list of neighbor ids
    neighbors: HashMap<u64, Vec<u64>>,
}

impl HnswLayer {
    /// Create a new empty layer.
    fn new() -> Self {
        Self {
            neighbors: HashMap::new(),
        }
    }

    /// Add a node to this layer.
    fn add_node(&mut self, id: u64) {
        self.neighbors.entry(id).or_default();
    }

    /// Check if a node exists in this layer.
    fn contains(&self, id: u64) -> bool {
        self.neighbors.contains_key(&id)
    }

    /// Get neighbors of a node.
    fn get_neighbors(&self, id: u64) -> &[u64] {
        self.neighbors.get(&id).map(|v| v.as_slice()).unwrap_or(&[])
    }

    /// Set neighbors for a node.
    fn set_neighbors(&mut self, id: u64, neighbors: Vec<u64>) {
        self.neighbors.insert(id, neighbors);
    }

    /// Add a bidirectional edge between two nodes.
    fn connect(&mut self, a: u64, b: u64, max_neighbors: usize) {
        // Add b to a's neighbors
        if let Some(neighbors) = self.neighbors.get_mut(&a) {
            if !neighbors.contains(&b) && neighbors.len() < max_neighbors {
                neighbors.push(b);
            }
        }
        // Add a to b's neighbors
        if let Some(neighbors) = self.neighbors.get_mut(&b) {
            if !neighbors.contains(&a) && neighbors.len() < max_neighbors {
                neighbors.push(a);
            }
        }
    }

    /// Remove a node and all its connections.
    fn remove_node(&mut self, id: u64) {
        // Remove from neighbors' lists
        if let Some(neighbors) = self.neighbors.remove(&id) {
            for neighbor_id in neighbors {
                if let Some(neighbor_neighbors) = self.neighbors.get_mut(&neighbor_id) {
                    neighbor_neighbors.retain(|&x| x != id);
                }
            }
        }
    }

    /// Number of nodes in this layer.
    fn len(&self) -> usize {
        self.neighbors.len()
    }

    /// Check if layer is empty.
    fn is_empty(&self) -> bool {
        self.neighbors.is_empty()
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// HNSW INDEX
// ═══════════════════════════════════════════════════════════════════════════════

/// HNSW index for approximate nearest neighbor search.
///
/// # Example
///
/// ```ignore
/// use lcpfs::vector::HnswIndex;
///
/// let mut index = HnswIndex::new(16, 200);
///
/// // Insert vectors
/// index.insert(1, &embedding1)?;
/// index.insert(2, &embedding2)?;
///
/// // Search
/// let results = index.search(&query, 10);
/// ```
#[derive(Clone)]
pub struct HnswIndex {
    /// Graph layers (layer 0 is the bottom/densest layer).
    layers: Vec<HnswLayer>,
    /// Entry point node ID (in the top layer).
    entry_point: Option<u64>,
    /// Maximum layer of the entry point.
    max_layer: usize,
    /// M parameter: max neighbors per node (layer > 0).
    m: usize,
    /// M_max0: max neighbors at layer 0 (typically 2*M).
    m_max0: usize,
    /// ef_construction: beam width during construction.
    ef_construction: usize,
    /// ef_search: beam width during search (can be adjusted at query time).
    ef_search: usize,
    /// Level multiplier for random layer selection.
    ml: f32,
    /// Distance metric.
    metric: DistanceMetric,
    /// Vector storage: id -> embedding.
    vectors: HashMap<u64, Vec<f32>>,
    /// Embedding dimensions.
    dimensions: usize,
    /// Simple RNG state for layer selection.
    rng_state: u64,
}

impl HnswIndex {
    /// Create a new HNSW index with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `m` - Maximum neighbors per node (default: 16)
    /// * `ef_construction` - Beam width during construction (default: 200)
    pub fn new(m: usize, ef_construction: usize) -> Self {
        let m = m.clamp(2, HNSW_MAX_NEIGHBORS);
        Self {
            layers: Vec::new(),
            entry_point: None,
            max_layer: 0,
            m,
            m_max0: m * 2,
            ef_construction,
            ef_search: 50,
            ml: 1.0 / libm::logf(m as f32),
            metric: DistanceMetric::Cosine,
            vectors: HashMap::new(),
            dimensions: 0,
            rng_state: 12345678, // Will be seeded properly
        }
    }

    /// Create with specific parameters.
    pub fn with_params(
        m: usize,
        ef_construction: usize,
        ef_search: usize,
        metric: DistanceMetric,
    ) -> Self {
        let mut index = Self::new(m, ef_construction);
        index.ef_search = ef_search;
        index.metric = metric;
        index
    }

    /// Set the search beam width.
    pub fn set_ef_search(&mut self, ef: usize) {
        self.ef_search = ef.max(1);
    }

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

    /// Check if the index is empty.
    pub fn is_empty(&self) -> bool {
        self.vectors.is_empty()
    }

    /// Get the embedding dimensions.
    pub fn dimensions(&self) -> usize {
        self.dimensions
    }

    /// Get the maximum layer.
    pub fn max_layer(&self) -> usize {
        self.max_layer
    }

    /// Seed the RNG with a value derived from the object ID.
    fn seed_rng(&mut self, id: u64) {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(id);
    }

    /// Generate a random float in [0, 1).
    fn random_f32(&mut self) -> f32 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1);
        (self.rng_state >> 33) as f32 / (1u64 << 31) as f32
    }

    /// Generate random layer for a new node.
    fn random_level(&mut self) -> usize {
        let r = self.random_f32();
        if r <= 0.0 {
            return 0;
        }
        let level = (-libm::logf(r) * self.ml) as usize;
        level.min(32) // Cap at reasonable max
    }

    /// Compute distance between two vectors by ID.
    fn distance(&self, id_a: u64, id_b: u64) -> f32 {
        let a = self.vectors.get(&id_a);
        let b = self.vectors.get(&id_b);
        match (a, b) {
            (Some(va), Some(vb)) => compute_distance(va, vb, self.metric),
            _ => f32::MAX,
        }
    }

    /// Compute distance between a query vector and a stored vector.
    fn distance_to_query(&self, query: &[f32], id: u64) -> f32 {
        match self.vectors.get(&id) {
            Some(v) => compute_distance(query, v, self.metric),
            None => f32::MAX,
        }
    }

    /// Ensure we have enough layers.
    fn ensure_layers(&mut self, level: usize) {
        while self.layers.len() <= level {
            self.layers.push(HnswLayer::new());
        }
    }

    /// Search layer for nearest neighbors (greedy search).
    ///
    /// Returns the nearest neighbor found starting from the entry point.
    fn search_layer_single(&self, query: &[f32], entry_point: u64, layer: usize) -> u64 {
        let mut current = entry_point;
        let mut current_dist = self.distance_to_query(query, current);

        loop {
            let mut changed = false;
            let neighbors = self.layers[layer].get_neighbors(current);

            for &neighbor in neighbors {
                let dist = self.distance_to_query(query, neighbor);
                if dist < current_dist {
                    current = neighbor;
                    current_dist = dist;
                    changed = true;
                }
            }

            if !changed {
                break;
            }
        }

        current
    }

    /// Search layer with beam search, returning ef-nearest candidates.
    fn search_layer(
        &self,
        query: &[f32],
        entry_points: &[u64],
        ef: usize,
        layer: usize,
    ) -> Vec<(u64, f32)> {
        // visited set
        let mut visited = hashbrown::HashSet::new();

        // candidates: min-heap by distance (nearest first)
        let mut candidates: BinaryHeap<MinCandidate> = BinaryHeap::new();

        // results: max-heap by distance (furthest first for eviction)
        let mut results: BinaryHeap<MaxCandidate> = BinaryHeap::new();

        // Initialize with entry points
        for &ep in entry_points {
            if visited.insert(ep) {
                let dist = self.distance_to_query(query, ep);
                candidates.push(MinCandidate {
                    id: ep,
                    distance: dist,
                });
                results.push(MaxCandidate {
                    id: ep,
                    distance: dist,
                });
            }
        }

        while let Some(MinCandidate {
            id: current,
            distance: current_dist,
        }) = candidates.pop()
        {
            // Get the furthest result
            let furthest_dist = results.peek().map(|r| r.distance).unwrap_or(f32::MAX);

            // If current candidate is further than furthest result, we're done
            if current_dist > furthest_dist {
                break;
            }

            // Explore neighbors
            let neighbors = self.layers[layer].get_neighbors(current);
            for &neighbor in neighbors {
                if visited.insert(neighbor) {
                    let dist = self.distance_to_query(query, neighbor);

                    // Add to candidates if closer than furthest result
                    if dist < furthest_dist || results.len() < ef {
                        candidates.push(MinCandidate {
                            id: neighbor,
                            distance: dist,
                        });
                        results.push(MaxCandidate {
                            id: neighbor,
                            distance: dist,
                        });

                        // Keep only ef best results
                        while results.len() > ef {
                            results.pop();
                        }
                    }
                }
            }
        }

        // Extract results sorted by distance (ascending)
        let mut result_vec: Vec<_> = results.into_iter().map(|c| (c.id, c.distance)).collect();
        result_vec.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        result_vec
    }

    /// Select neighbors using simple heuristic (keep nearest M).
    fn select_neighbors_simple(&self, candidates: &[(u64, f32)], m: usize) -> Vec<u64> {
        candidates.iter().take(m).map(|(id, _)| *id).collect()
    }

    /// Select neighbors using heuristic that considers diversity.
    ///
    /// This prevents the graph from becoming too clustered by preferring
    /// neighbors that provide different "directions" from the node.
    fn select_neighbors_heuristic(
        &self,
        query_id: u64,
        candidates: &[(u64, f32)],
        m: usize,
        layer: usize,
        extend_candidates: bool,
    ) -> Vec<u64> {
        if candidates.len() <= m {
            return candidates.iter().map(|(id, _)| *id).collect();
        }

        let mut working_candidates = candidates.to_vec();

        // Optionally extend with neighbors of candidates
        if extend_candidates {
            let mut extended = hashbrown::HashSet::new();
            for (id, _) in &working_candidates {
                extended.insert(*id);
            }

            for (id, _) in candidates {
                for &neighbor in self.layers[layer].get_neighbors(*id) {
                    if neighbor != query_id && extended.insert(neighbor) {
                        let dist = self.distance(query_id, neighbor);
                        working_candidates.push((neighbor, dist));
                    }
                }
            }

            working_candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        }

        // Select neighbors preferring those not too close to already selected
        let mut selected = Vec::with_capacity(m);
        let mut selected_set = hashbrown::HashSet::new();

        for (id, dist) in &working_candidates {
            if selected.len() >= m {
                break;
            }

            // Check if this candidate is closer to query than to any selected neighbor
            let mut good = true;
            for &sel_id in &selected {
                let dist_to_selected = self.distance(*id, sel_id);
                if dist_to_selected < *dist {
                    good = false;
                    break;
                }
            }

            if good && selected_set.insert(*id) {
                selected.push(*id);
            }
        }

        // If we couldn't select enough diverse neighbors, fill with nearest
        if selected.len() < m {
            for (id, _) in &working_candidates {
                if selected.len() >= m {
                    break;
                }
                if selected_set.insert(*id) {
                    selected.push(*id);
                }
            }
        }

        selected
    }

    /// Insert a vector into the index.
    ///
    /// # Arguments
    ///
    /// * `id` - Unique identifier for this vector
    /// * `embedding` - The vector embedding
    ///
    /// # Returns
    ///
    /// `Ok(())` on success, or an error if dimensions don't match.
    pub fn insert(&mut self, id: u64, embedding: &[f32]) -> Result<(), VectorError> {
        // Check dimensions
        if self.dimensions == 0 {
            self.dimensions = embedding.len();
        } else if embedding.len() != self.dimensions {
            return Err(VectorError::DimensionMismatch {
                expected: self.dimensions,
                actual: embedding.len(),
            });
        }

        // Store the vector
        self.vectors.insert(id, embedding.to_vec());

        // Seed RNG and generate random level
        self.seed_rng(id);
        let level = self.random_level();

        // Ensure we have enough layers
        self.ensure_layers(level);

        // Handle first insertion
        if self.entry_point.is_none() {
            self.entry_point = Some(id);
            self.max_layer = level;
            for l in 0..=level {
                self.layers[l].add_node(id);
            }
            return Ok(());
        }

        let entry_point = self.entry_point.unwrap();

        // Phase 1: Greedy search from top to level+1
        let mut current = entry_point;
        for l in (level + 1..=self.max_layer).rev() {
            current = self.search_layer_single(embedding, current, l);
        }

        // Phase 2: Insert at each layer from level down to 0
        for l in (0..=level.min(self.max_layer)).rev() {
            // Add node to this layer
            self.layers[l].add_node(id);

            // Find nearest neighbors
            let candidates = self.search_layer(embedding, &[current], self.ef_construction, l);

            // Select neighbors
            let m = if l == 0 { self.m_max0 } else { self.m };
            let neighbors = self.select_neighbors_heuristic(id, &candidates, m, l, true);

            // Connect to neighbors
            self.layers[l].set_neighbors(id, neighbors.clone());

            // Add reverse edges and prune if necessary
            for &neighbor in &neighbors {
                let mut neighbor_neighbors: Vec<u64> =
                    self.layers[l].get_neighbors(neighbor).to_vec();

                if !neighbor_neighbors.contains(&id) {
                    neighbor_neighbors.push(id);

                    // Prune if too many neighbors
                    if neighbor_neighbors.len() > m {
                        // Get distances and select best
                        let mut with_dist: Vec<_> = neighbor_neighbors
                            .iter()
                            .map(|&n| (n, self.distance(neighbor, n)))
                            .collect();
                        with_dist.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
                        neighbor_neighbors =
                            with_dist.into_iter().take(m).map(|(n, _)| n).collect();
                    }

                    self.layers[l].set_neighbors(neighbor, neighbor_neighbors);
                }
            }

            // Update current for next layer
            if !candidates.is_empty() {
                current = candidates[0].0;
            }
        }

        // Update entry point if new node is at higher level
        if level > self.max_layer {
            self.entry_point = Some(id);
            self.max_layer = level;
        }

        Ok(())
    }

    /// Search for k nearest neighbors.
    ///
    /// # Arguments
    ///
    /// * `query` - Query vector
    /// * `k` - Number of neighbors to return
    ///
    /// # Returns
    ///
    /// Vector of search results sorted by distance (ascending).
    pub fn search(&self, query: &[f32], k: usize) -> Vec<VectorSearchResult> {
        if self.is_empty() || query.len() != self.dimensions {
            return Vec::new();
        }

        let entry_point = match self.entry_point {
            Some(ep) => ep,
            None => return Vec::new(),
        };

        // Phase 1: Greedy search from top to layer 1
        let mut current = entry_point;
        for l in (1..=self.max_layer).rev() {
            current = self.search_layer_single(query, current, l);
        }

        // Phase 2: Search layer 0 with ef candidates
        let ef = self.ef_search.max(k);
        let candidates = self.search_layer(query, &[current], ef, 0);

        // Return top-k results
        candidates
            .into_iter()
            .take(k)
            .map(|(id, distance)| {
                // Convert distance to similarity score
                let score = match self.metric {
                    DistanceMetric::Cosine => 1.0 - distance, // cosine distance -> similarity
                    DistanceMetric::DotProduct => -distance,  // negative dot -> positive
                    _ => 1.0 / (1.0 + distance),              // generic: closer = higher
                };
                VectorSearchResult::new(id, score, distance)
            })
            .collect()
    }

    /// Search with custom ef parameter.
    pub fn search_with_ef(&self, query: &[f32], k: usize, ef: usize) -> Vec<VectorSearchResult> {
        if self.is_empty() || query.len() != self.dimensions {
            return Vec::new();
        }

        let entry_point = match self.entry_point {
            Some(ep) => ep,
            None => return Vec::new(),
        };

        let mut current = entry_point;
        for l in (1..=self.max_layer).rev() {
            current = self.search_layer_single(query, current, l);
        }

        let candidates = self.search_layer(query, &[current], ef.max(k), 0);

        candidates
            .into_iter()
            .take(k)
            .map(|(id, distance)| {
                let score = match self.metric {
                    DistanceMetric::Cosine => 1.0 - distance,
                    DistanceMetric::DotProduct => -distance,
                    _ => 1.0 / (1.0 + distance),
                };
                VectorSearchResult::new(id, score, distance)
            })
            .collect()
    }

    /// Delete a vector from the index.
    ///
    /// Note: HNSW deletion is complex and can degrade search quality.
    /// For high deletion rates, consider rebuilding the index periodically.
    pub fn delete(&mut self, id: u64) -> Result<(), VectorError> {
        if !self.vectors.contains_key(&id) {
            return Err(VectorError::ObjectNotFound(id));
        }

        // Remove from all layers
        for layer in &mut self.layers {
            layer.remove_node(id);
        }

        // Remove vector data
        self.vectors.remove(&id);

        // If we deleted the entry point, find a new one
        if self.entry_point == Some(id) {
            self.entry_point = None;
            self.max_layer = 0;

            // Find new entry point (node at highest layer)
            for (l, layer) in self.layers.iter().enumerate().rev() {
                if !layer.is_empty() {
                    // Get any node from this layer
                    if let Some(&first) = layer.neighbors.keys().next() {
                        self.entry_point = Some(first);
                        self.max_layer = l;
                        break;
                    }
                }
            }
        }

        Ok(())
    }

    /// Get a stored vector by ID.
    pub fn get_vector(&self, id: u64) -> Option<&[f32]> {
        self.vectors.get(&id).map(|v| v.as_slice())
    }

    /// Check if a vector exists in the index.
    pub fn contains(&self, id: u64) -> bool {
        self.vectors.contains_key(&id)
    }

    /// Get all vector IDs in the index.
    pub fn get_ids(&self) -> Vec<u64> {
        self.vectors.keys().copied().collect()
    }

    /// Get statistics about the index.
    pub fn stats(&self) -> HnswStats {
        let mut layer_sizes = Vec::new();
        let mut total_edges = 0;

        for layer in &self.layers {
            layer_sizes.push(layer.len());
            for neighbors in layer.neighbors.values() {
                total_edges += neighbors.len();
            }
        }

        HnswStats {
            vector_count: self.vectors.len(),
            dimensions: self.dimensions,
            layer_count: self.layers.len(),
            layer_sizes,
            total_edges: total_edges / 2, // Edges are bidirectional
            m: self.m,
            ef_construction: self.ef_construction,
            ef_search: self.ef_search,
            metric: self.metric,
            entry_point: self.entry_point,
        }
    }

    /// Get the distance metric used by this index.
    pub fn metric(&self) -> DistanceMetric {
        self.metric
    }

    /// Get the entry point node ID.
    pub fn entry_point(&self) -> Option<u64> {
        self.entry_point
    }
}

/// Statistics about an HNSW index.
#[derive(Debug, Clone)]
pub struct HnswStats {
    /// Total number of vectors.
    pub vector_count: usize,
    /// Embedding dimensions.
    pub dimensions: usize,
    /// Number of layers.
    pub layer_count: usize,
    /// Number of nodes in each layer.
    pub layer_sizes: Vec<usize>,
    /// Total number of edges in the graph.
    pub total_edges: usize,
    /// M parameter.
    pub m: usize,
    /// ef_construction parameter.
    pub ef_construction: usize,
    /// ef_search parameter.
    pub ef_search: usize,
    /// Distance metric used by this index.
    pub metric: DistanceMetric,
    /// Entry point node ID (None if index is empty).
    pub entry_point: Option<u64>,
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

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

    fn random_vector(dim: usize, seed: u64) -> Vec<f32> {
        let mut rng = seed;
        (0..dim)
            .map(|_| {
                rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
                ((rng >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0
            })
            .collect()
    }

    fn normalize(v: &mut [f32]) {
        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for x in v.iter_mut() {
                *x /= norm;
            }
        }
    }

    #[test]
    fn test_hnsw_empty() {
        let index = HnswIndex::new(16, 200);
        assert!(index.is_empty());
        assert_eq!(index.len(), 0);

        let query = vec![1.0, 0.0, 0.0];
        let results = index.search(&query, 10);
        assert!(results.is_empty());
    }

    #[test]
    fn test_hnsw_single_insert() {
        let mut index = HnswIndex::new(16, 200);
        let embedding = vec![1.0, 0.0, 0.0];

        index.insert(1, &embedding).unwrap();

        assert_eq!(index.len(), 1);
        assert!(!index.is_empty());
        assert!(index.contains(1));
        assert!(!index.contains(2));
    }

    #[test]
    fn test_hnsw_search_exact() {
        let mut index = HnswIndex::new(16, 200);

        // Insert a few vectors
        index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0]).unwrap();
        index.insert(3, &[0.0, 0.0, 1.0]).unwrap();

        // Search for exact match
        let results = index.search(&[1.0, 0.0, 0.0], 1);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].object_id, 1);
        assert!(results[0].distance < 0.01);
    }

    #[test]
    fn test_hnsw_search_nearest() {
        let mut index = HnswIndex::new(16, 200);

        // Insert vectors
        index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.9, 0.1, 0.0]).unwrap(); // Close to 1
        index.insert(3, &[0.0, 1.0, 0.0]).unwrap();
        index.insert(4, &[0.0, 0.0, 1.0]).unwrap();

        // Search for something close to [1, 0, 0]
        let results = index.search(&[0.95, 0.05, 0.0], 2);
        assert!(results.len() >= 2);

        // Results should be ordered by distance (ascending) or score (descending)
        // First result should be id 1 or 2 (both are close)
        assert!(results[0].object_id == 1 || results[0].object_id == 2);
    }

    #[test]
    fn test_hnsw_delete() {
        let mut index = HnswIndex::new(16, 200);

        index.insert(1, &[1.0, 0.0, 0.0]).unwrap();
        index.insert(2, &[0.0, 1.0, 0.0]).unwrap();
        index.insert(3, &[0.0, 0.0, 1.0]).unwrap();

        assert_eq!(index.len(), 3);

        index.delete(2).unwrap();
        assert_eq!(index.len(), 2);
        assert!(!index.contains(2));

        // Search should not return deleted vector
        let results = index.search(&[0.0, 1.0, 0.0], 10);
        for r in &results {
            assert_ne!(r.object_id, 2);
        }
    }

    #[test]
    fn test_hnsw_dimension_mismatch() {
        let mut index = HnswIndex::new(16, 200);

        index.insert(1, &[1.0, 0.0, 0.0]).unwrap();

        let result = index.insert(2, &[1.0, 0.0]); // Wrong dimensions
        assert!(matches!(result, Err(VectorError::DimensionMismatch { .. })));
    }

    #[test]
    fn test_hnsw_larger_dataset() {
        let mut index = HnswIndex::new(16, 200);
        let dim = 64;
        let n = 100;

        // Insert random vectors
        for i in 0..n {
            let mut v = random_vector(dim, i as u64);
            normalize(&mut v);
            index.insert(i as u64, &v).unwrap();
        }

        assert_eq!(index.len(), n);

        // Search should return results
        let query = random_vector(dim, 999);
        let results = index.search(&query, 10);
        assert!(results.len() <= 10);
        assert!(!results.is_empty());

        // Check stats
        let stats = index.stats();
        assert_eq!(stats.vector_count, n);
        assert_eq!(stats.dimensions, dim);
    }

    #[test]
    fn test_hnsw_recall() {
        // Test that HNSW achieves good recall on a moderate dataset
        let mut index = HnswIndex::new(16, 200);
        index.set_ef_search(100);

        let dim = 32;
        let n = 200;

        // Generate and insert vectors
        let mut vectors: Vec<Vec<f32>> = Vec::new();
        for i in 0..n {
            let mut v = random_vector(dim, i as u64);
            normalize(&mut v);
            vectors.push(v.clone());
            index.insert(i as u64, &v).unwrap();
        }

        // Test recall on several queries
        let mut total_recall = 0.0;
        let num_queries = 10;
        let k = 10;

        for q in 0..num_queries {
            let query = &vectors[q * 10]; // Use some vectors as queries

            // Get approximate results from HNSW
            let approx_results = index.search(query, k);
            let approx_ids: hashbrown::HashSet<_> =
                approx_results.iter().map(|r| r.object_id).collect();

            // Compute exact nearest neighbors (brute force)
            let mut exact: Vec<_> = (0..n as u64)
                .map(|i| {
                    let dist = cosine_distance(query, &vectors[i as usize]);
                    (i, dist)
                })
                .collect();
            exact.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

            // Count how many of top-k exact are in approximate results
            let mut hits = 0;
            for (id, _) in exact.iter().take(k) {
                if approx_ids.contains(id) {
                    hits += 1;
                }
            }

            total_recall += hits as f32 / k as f32;
        }

        let avg_recall = total_recall / num_queries as f32;

        // With ef_search=100 and moderate dataset, expect >90% recall
        assert!(
            avg_recall > 0.8,
            "Average recall {} is too low (expected > 0.8)",
            avg_recall
        );
    }

    #[test]
    fn test_hnsw_get_vector() {
        let mut index = HnswIndex::new(16, 200);
        let embedding = vec![1.0, 2.0, 3.0];

        index.insert(42, &embedding).unwrap();

        let retrieved = index.get_vector(42).unwrap();
        assert_eq!(retrieved, &[1.0, 2.0, 3.0]);

        assert!(index.get_vector(999).is_none());
    }
}