ruv-neural-memory 0.1.0

//! Simplified HNSW (Hierarchical Navigable Small World) index for approximate
//! nearest neighbor search on embedding vectors.

use std::collections::{BinaryHeap, HashSet};
use std::cmp::Ordering;

/// A scored neighbor for use in the priority queue.
#[derive(Debug, Clone)]
struct ScoredNode {
    id: usize,
    distance: f64,
}

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

impl Eq for ScoredNode {}

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

impl Ord for ScoredNode {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse ordering for min-heap behavior
        other
            .distance
            .partial_cmp(&self.distance)
            .unwrap_or(Ordering::Equal)
    }
}

/// Max-heap scored node (furthest first).
#[derive(Debug, Clone)]
struct FurthestNode {
    id: usize,
    distance: f64,
}

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

impl Eq for FurthestNode {}

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

impl Ord for FurthestNode {
    fn cmp(&self, other: &Self) -> Ordering {
        self.distance
            .partial_cmp(&other.distance)
            .unwrap_or(Ordering::Equal)
    }
}

/// Hierarchical Navigable Small World graph for approximate nearest neighbor search.
///
/// This is a simplified single-layer HNSW implementation suitable for moderate-scale
/// embedding stores (up to ~100k vectors).
pub struct HnswIndex {
    /// Adjacency list per layer: layers[layer][node] = [(neighbor_id, distance)]
    layers: Vec<Vec<Vec<(usize, f64)>>>,
    /// Entry point node for search.
    entry_point: usize,
    /// Maximum layer index currently in the graph.
    max_layer: usize,
    /// Number of neighbors to consider during construction.
    ef_construction: usize,
    /// Maximum number of connections per node per layer.
    m: usize,
    /// Stored embedding vectors.
    embeddings: Vec<Vec<f64>>,
}

impl HnswIndex {
    /// Create a new empty HNSW index.
    ///
    /// - `m`: maximum connections per node per layer (typical: 16)
    /// - `ef_construction`: search width during construction (typical: 200)
    pub fn new(m: usize, ef_construction: usize) -> Self {
        Self {
            layers: vec![Vec::new()], // Start with layer 0
            entry_point: 0,
            max_layer: 0,
            ef_construction,
            m,
            embeddings: Vec::new(),
        }
    }

    /// Insert a vector and return its index.
    pub fn insert(&mut self, vector: &[f64]) -> usize {
        let id = self.embeddings.len();
        self.embeddings.push(vector.to_vec());

        let insert_layer = self.select_layer();

        // Ensure we have enough layers
        while self.layers.len() <= insert_layer {
            self.layers.push(Vec::new());
        }

        // Add empty adjacency lists for this node in all layers up to insert_layer
        for layer in 0..=insert_layer {
            while self.layers[layer].len() <= id {
                self.layers[layer].push(Vec::new());
            }
        }

        // Also ensure layer 0 has an entry for this node
        while self.layers[0].len() <= id {
            self.layers[0].push(Vec::new());
        }

        if id == 0 {
            // First node, just set as entry point
            self.entry_point = 0;
            self.max_layer = insert_layer;
            return id;
        }

        // Greedy search from top layer down to insert_layer+1
        let mut current_entry = self.entry_point;
        for layer in (insert_layer + 1..=self.max_layer).rev() {
            if layer < self.layers.len() {
                let neighbors = self.search_layer(vector, current_entry, 1, layer);
                if let Some((nearest, _)) = neighbors.first() {
                    current_entry = *nearest;
                }
            }
        }

        // Insert into layers from insert_layer down to 0
        for layer in (0..=insert_layer.min(self.max_layer)).rev() {
            let neighbors =
                self.search_layer(vector, current_entry, self.ef_construction, layer);

            // Select up to m neighbors
            let selected: Vec<(usize, f64)> =
                neighbors.into_iter().take(self.m).collect();

            // Ensure adjacency list exists for this node at this layer
            while self.layers[layer].len() <= id {
                self.layers[layer].push(Vec::new());
            }

            // Add bidirectional connections
            for &(neighbor_id, dist) in &selected {
                self.layers[layer][id].push((neighbor_id, dist));

                while self.layers[layer].len() <= neighbor_id {
                    self.layers[layer].push(Vec::new());
                }
                self.layers[layer][neighbor_id].push((id, dist));

                // Prune if over capacity
                if self.layers[layer][neighbor_id].len() > self.m * 2 {
                    self.layers[layer][neighbor_id]
                        .sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
                    self.layers[layer][neighbor_id].truncate(self.m * 2);
                }
            }

            if let Some((nearest, _)) = selected.first() {
                current_entry = *nearest;
            }
        }

        if insert_layer > self.max_layer {
            self.max_layer = insert_layer;
            self.entry_point = id;
        }

        id
    }

    /// Search for the k nearest neighbors of `query`.
    ///
    /// - `k`: number of nearest neighbors to return
    /// - `ef`: search width (larger = more accurate, slower; typical: 50-200)
    ///
    /// Returns (index, distance) pairs sorted by ascending distance.
    pub fn search(&self, query: &[f64], k: usize, ef: usize) -> Vec<(usize, f64)> {
        if self.embeddings.is_empty() {
            return Vec::new();
        }

        // Bounds-check the entry point
        if self.entry_point >= self.embeddings.len() {
            return Vec::new();
        }

        let mut current_entry = self.entry_point;

        // Greedy search from top layer down to layer 1
        for layer in (1..=self.max_layer).rev() {
            if layer < self.layers.len() {
                let neighbors = self.search_layer(query, current_entry, 1, layer);
                if let Some((nearest, _)) = neighbors.first() {
                    current_entry = *nearest;
                }
            }
        }

        // Search layer 0 with ef candidates
        let mut results = self.search_layer(query, current_entry, ef.max(k), 0);
        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
        results.truncate(k);
        results
    }

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

    /// Returns true if the index has no vectors.
    pub fn is_empty(&self) -> bool {
        self.embeddings.is_empty()
    }

    /// Euclidean distance between two vectors.
    fn distance(a: &[f64], b: &[f64]) -> f64 {
        a.iter()
            .zip(b.iter())
            .map(|(x, y)| (x - y) * (x - y))
            .sum::<f64>()
            .sqrt()
    }

    /// Select a random layer for insertion using an exponential distribution.
    fn select_layer(&self) -> usize {
        // Deterministic level assignment based on node count for reproducibility.
        // Uses a simple hash-like scheme: most nodes go to layer 0.
        let n = self.embeddings.len();
        let ml = 1.0 / (self.m as f64).ln();
        // Use a simple deterministic pseudo-random based on n
        let hash = ((n.wrapping_mul(2654435761)) >> 16) as f64 / 65536.0;
        let level = (-hash.ln() * ml).floor() as usize;
        level.min(4) // Cap at 4 layers
    }

    /// Search a single layer starting from `entry`, returning `ef` nearest candidates.
    fn search_layer(
        &self,
        query: &[f64],
        entry: usize,
        ef: usize,
        layer: usize,
    ) -> Vec<(usize, f64)> {
        if layer >= self.layers.len() {
            return Vec::new();
        }

        // Bounds-check entry against embeddings
        if entry >= self.embeddings.len() {
            return Vec::new();
        }

        let mut visited = HashSet::new();
        let entry_dist = Self::distance(query, &self.embeddings[entry]);

        // Candidates: min-heap (closest first)
        let mut candidates = BinaryHeap::new();
        candidates.push(ScoredNode {
            id: entry,
            distance: entry_dist,
        });

        // Results: max-heap (furthest first, for pruning)
        let mut results = BinaryHeap::new();
        results.push(FurthestNode {
            id: entry,
            distance: entry_dist,
        });

        visited.insert(entry);

        while let Some(ScoredNode { id: current, distance: current_dist }) = candidates.pop() {
            // If current candidate is further than the worst result and we have enough, stop
            if let Some(worst) = results.peek() {
                if current_dist > worst.distance && results.len() >= ef {
                    break;
                }
            }

            // Explore neighbors
            if current < self.layers[layer].len() {
                for &(neighbor, _) in &self.layers[layer][current] {
                    if neighbor < self.embeddings.len() && visited.insert(neighbor) {
                        let dist = Self::distance(query, &self.embeddings[neighbor]);

                        let should_add = results.len() < ef
                            || results
                                .peek()
                                .map(|w| dist < w.distance)
                                .unwrap_or(true);

                        if should_add {
                            candidates.push(ScoredNode {
                                id: neighbor,
                                distance: dist,
                            });
                            results.push(FurthestNode {
                                id: neighbor,
                                distance: dist,
                            });

                            if results.len() > ef {
                                results.pop();
                            }
                        }
                    }
                }
            }
        }

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

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

    #[test]
    fn insert_and_search_basic() {
        let mut index = HnswIndex::new(4, 20);
        index.insert(&[0.0, 0.0]);
        index.insert(&[1.0, 0.0]);
        index.insert(&[0.0, 1.0]);
        index.insert(&[10.0, 10.0]);

        let results = index.search(&[0.1, 0.1], 2, 10);
        assert_eq!(results.len(), 2);
        // Closest should be [0,0]
        assert_eq!(results[0].0, 0);
    }

    #[test]
    fn empty_index_returns_empty() {
        let index = HnswIndex::new(4, 20);
        let results = index.search(&[1.0, 2.0], 5, 10);
        assert!(results.is_empty());
    }

    #[test]
    fn single_element() {
        let mut index = HnswIndex::new(4, 20);
        index.insert(&[5.0, 5.0]);

        let results = index.search(&[0.0, 0.0], 1, 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 0);
    }

    #[test]
    fn hnsw_recall_vs_brute_force() {
        use rand::Rng;

        let mut rng = rand::thread_rng();
        let dim = 8;
        let n = 200;
        let k = 10;

        let mut index = HnswIndex::new(16, 100);
        let mut vectors: Vec<Vec<f64>> = Vec::new();

        for _ in 0..n {
            let v: Vec<f64> = (0..dim).map(|_| rng.gen_range(-1.0..1.0)).collect();
            index.insert(&v);
            vectors.push(v);
        }

        // Run multiple queries and check average recall
        let num_queries = 20;
        let mut total_recall = 0.0;

        for _ in 0..num_queries {
            let query: Vec<f64> = (0..dim).map(|_| rng.gen_range(-1.0..1.0)).collect();

            // Brute force ground truth
            let mut bf_distances: Vec<(usize, f64)> = vectors
                .iter()
                .enumerate()
                .map(|(i, v)| (i, HnswIndex::distance(&query, v)))
                .collect();
            bf_distances
                .sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
            let bf_top_k: Vec<usize> = bf_distances.iter().take(k).map(|(i, _)| *i).collect();

            // HNSW search
            let hnsw_results = index.search(&query, k, 50);
            let hnsw_top_k: Vec<usize> = hnsw_results.iter().map(|(i, _)| *i).collect();

            // Compute recall
            let hits = hnsw_top_k
                .iter()
                .filter(|id| bf_top_k.contains(id))
                .count();
            total_recall += hits as f64 / k as f64;
        }

        let avg_recall = total_recall / num_queries as f64;
        assert!(
            avg_recall > 0.9,
            "HNSW recall {} should be > 0.9",
            avg_recall
        );
    }

    #[test]
    fn distance_is_euclidean() {
        let d = HnswIndex::distance(&[0.0, 0.0], &[3.0, 4.0]);
        assert!((d - 5.0).abs() < 1e-10);
    }
}