vicinity 0.5.0

//! HNSW graph construction algorithm.

use crate::distance;
use crate::hnsw::graph::{HNSWIndex, Layer};
use crate::hnsw::search::greedy_search_layer;
use crate::RetrieveError;
use smallvec::SmallVec;

/// Return a plain function pointer for the configured distance metric.
#[inline(always)]
fn dist_fn_for(index: &HNSWIndex) -> fn(&[f32], &[f32]) -> f32 {
    use crate::distance::DistanceMetric;
    match index.params.metric {
        DistanceMetric::L2 => distance::l2_distance,
        DistanceMetric::Cosine => distance::cosine_distance_normalized,
        DistanceMetric::Angular => distance::angular_distance,
        DistanceMetric::InnerProduct => distance::inner_product_distance,
    }
}

/// Select neighbors using RND (Relative Neighborhood Diversification).
///
/// Criterion: include candidate \(X_j\) if it is closer to the query than it is to
/// every already-selected neighbor \(X_i\) (i.e. dist(q, j) < dist(i, j) for all selected i).
fn select_neighbors_rnd(
    _query_vector: &[f32],
    candidates: &mut [(u32, f32)],
    m: usize,
    vectors: &[f32],
    dimension: usize,
    dist_fn: fn(&[f32], &[f32]) -> f32,
) -> Vec<u32> {
    if candidates.is_empty() {
        return Vec::new();
    }

    // Sort in place (caller passes owned/mutable slice)
    candidates.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));

    let mut selected = Vec::with_capacity(m.min(candidates.len()));
    // Collect rejected candidates in distance order for O(1) backfill
    let mut rejected = Vec::new();

    // Start with closest candidate
    selected.push(candidates[0].0);

    // RND: Add candidate X_j if dist(X_q, X_j) < dist(X_i, X_j) for all selected neighbors X_i
    for &(candidate_id, query_to_candidate_dist) in candidates.iter().skip(1) {
        if selected.len() >= m {
            break;
        }

        let candidate_vec = get_vector(vectors, dimension, candidate_id as usize);
        let mut can_add = true;

        for &selected_id in &selected {
            let selected_vec = get_vector(vectors, dimension, selected_id as usize);
            let inter_distance = dist_fn(selected_vec, candidate_vec);

            if query_to_candidate_dist >= inter_distance {
                can_add = false;
                break;
            }
        }

        if can_add {
            selected.push(candidate_id);
        } else {
            rejected.push(candidate_id);
        }
    }

    // Backfill from rejected candidates (already in distance order)
    let mut ri = 0;
    while selected.len() < m && ri < rejected.len() {
        selected.push(rejected[ri]);
        ri += 1;
    }

    selected
}

/// Select neighbors using MOND (Maximum-Oriented Neighborhood Diversification).
///
/// Maximizes angles between neighbors. Formula: ∠(X_j X_q X_i) > θ for all selected X_i.
/// Second-best ND strategy with moderate pruning (2-4%).
fn select_neighbors_mond(
    query_vector: &[f32],
    candidates: &mut [(u32, f32)],
    m: usize,
    vectors: &[f32],
    dimension: usize,
    min_angle_degrees: f32,
) -> Vec<u32> {
    if candidates.is_empty() {
        return Vec::new();
    }

    let min_angle_rad = min_angle_degrees.to_radians();
    let min_cos = min_angle_rad.cos();

    // Sort in place
    candidates.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));

    let mut selected = Vec::with_capacity(m.min(candidates.len()));
    let mut rejected = Vec::new();

    // Start with closest candidate
    selected.push(candidates[0].0);

    // Precompute dot(query, query) once
    use crate::simd;
    let dot_qq = simd::dot(query_vector, query_vector);

    // MOND: Add candidate if angle with all selected neighbors > min_angle
    for &(candidate_id, _) in candidates.iter().skip(1) {
        if selected.len() >= m {
            break;
        }

        let candidate_vec = get_vector(vectors, dimension, candidate_id as usize);
        let mut can_add = true;

        let dot_cq = simd::dot(candidate_vec, query_vector);
        let dot_cc_self = simd::dot(candidate_vec, candidate_vec);

        for &selected_id in &selected {
            let selected_vec = get_vector(vectors, dimension, selected_id as usize);

            let dot_cc = simd::dot(candidate_vec, selected_vec);
            let dot_sq = simd::dot(selected_vec, query_vector);
            let dot_qc_qs = dot_cc - dot_cq - dot_sq + dot_qq;

            let norm_c_sq = dot_cc_self + dot_qq - 2.0 * dot_cq;
            let norm_s_sq = simd::dot(selected_vec, selected_vec) + dot_qq - 2.0 * dot_sq;

            if norm_c_sq > 0.0 && norm_s_sq > 0.0 {
                let norm_c = norm_c_sq.sqrt();
                let norm_s = norm_s_sq.sqrt();
                let cos_angle = dot_qc_qs / (norm_c * norm_s);
                if cos_angle >= min_cos {
                    can_add = false;
                    break;
                }
            }
        }

        if can_add {
            selected.push(candidate_id);
        } else {
            rejected.push(candidate_id);
        }
    }

    // Backfill from rejected (already in distance order)
    let mut ri = 0;
    while selected.len() < m && ri < rejected.len() {
        selected.push(rejected[ri]);
        ri += 1;
    }

    selected
}

/// Select neighbors using RRND (Relaxed Relative Neighborhood Diversification).
///
/// Formula: dist(X_q, X_j) < α · dist(X_i, X_j) with α ≥ 1.5.
/// Less effective than RND, creates larger graphs.
fn select_neighbors_rrnd(
    _query_vector: &[f32],
    candidates: &mut [(u32, f32)],
    m: usize,
    vectors: &[f32],
    dimension: usize,
    alpha: f32,
    dist_fn: fn(&[f32], &[f32]) -> f32,
) -> Vec<u32> {
    if candidates.is_empty() {
        return Vec::new();
    }

    // Sort in place
    candidates.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));

    let mut selected = Vec::with_capacity(m.min(candidates.len()));
    let mut rejected = Vec::new();

    // Start with closest candidate
    selected.push(candidates[0].0);

    // RRND: Add candidate X_j if dist(X_q, X_j) < α · dist(X_i, X_j) for all selected X_i
    for &(candidate_id, query_to_candidate_dist) in candidates.iter().skip(1) {
        if selected.len() >= m {
            break;
        }

        let candidate_vec = get_vector(vectors, dimension, candidate_id as usize);
        let mut can_add = true;

        for &selected_id in &selected {
            let selected_vec = get_vector(vectors, dimension, selected_id as usize);
            let inter_distance = dist_fn(selected_vec, candidate_vec);

            if query_to_candidate_dist >= alpha * inter_distance {
                can_add = false;
                break;
            }
        }

        if can_add {
            selected.push(candidate_id);
        } else {
            rejected.push(candidate_id);
        }
    }

    // Backfill from rejected (already in distance order)
    let mut ri = 0;
    while selected.len() < m && ri < rejected.len() {
        selected.push(rejected[ri]);
        ri += 1;
    }

    selected
}

/// Select neighbors based on configured diversification strategy.
///
/// `candidates` is sorted in place to avoid allocation.
pub fn select_neighbors(
    query_vector: &[f32],
    candidates: &mut [(u32, f32)],
    m: usize,
    vectors: &[f32],
    dimension: usize,
    strategy: &crate::hnsw::graph::NeighborhoodDiversification,
    dist_fn: fn(&[f32], &[f32]) -> f32,
) -> Vec<u32> {
    match strategy {
        crate::hnsw::graph::NeighborhoodDiversification::RelativeNeighborhood => {
            select_neighbors_rnd(query_vector, candidates, m, vectors, dimension, dist_fn)
        }
        crate::hnsw::graph::NeighborhoodDiversification::MaximumOriented { min_angle_degrees } => {
            select_neighbors_mond(
                query_vector,
                candidates,
                m,
                vectors,
                dimension,
                *min_angle_degrees,
            )
        }
        crate::hnsw::graph::NeighborhoodDiversification::RelaxedRelative { alpha } => {
            select_neighbors_rrnd(
                query_vector,
                candidates,
                m,
                vectors,
                dimension,
                *alpha,
                dist_fn,
            )
        }
    }
}

/// Get vector from SoA storage.
#[inline]
pub fn get_vector(vectors: &[f32], dimension: usize, idx: usize) -> &[f32] {
    let start = idx * dimension;
    let end = start + dimension;
    &vectors[start..end]
}

/// Construct HNSW graph layers.
///
/// Implements the insertion algorithm from the HNSW paper (Malkov & Yashunin, 2018).
///
/// Key insight: When descending through layers, we use the closest node found
/// in the layer above as the entry point for the next layer. This ensures
/// we start searching from a good position, not an arbitrary node.
pub fn construct_graph(index: &mut HNSWIndex) -> Result<(), RetrieveError> {
    if index.num_vectors == 0 {
        return Err(RetrieveError::EmptyIndex);
    }

    // Find maximum layer
    let max_layer = index.layer_assignments.iter().max().copied().unwrap_or(0) as usize;

    // Initialize layers with uncompressed storage
    index.layers = (0..=max_layer)
        .map(|_| Layer::new_uncompressed(vec![SmallVec::new(); index.num_vectors]))
        .collect();

    // Capture the distance function once so later mutable borrows of `index` don't conflict.
    let dist_fn = dist_fn_for(index);

    // Global entry point for the already-inserted subgraph.
    //
    // IMPORTANT: we are doing an *offline* build from stored vectors, but we still must
    // respect insertion order: the entry point must be a node that has already been
    // inserted (otherwise early nodes route through "future" nodes that have no edges yet).
    let mut global_entry_point = 0u32;
    let mut global_entry_layer = index.layer_assignments[0];

    // Insert each vector into the graph
    for current_id in 0..index.num_vectors {
        let current_layer = index.layer_assignments[current_id] as usize;
        let current_vector = get_vector(&index.vectors, index.dimension, current_id);

        // First node initializes the entry point.
        if current_id == 0 {
            global_entry_point = 0;
            global_entry_layer = index.layer_assignments[0];
            continue;
        }

        // Track closest node found while descending through layers.
        // We propagate the best entry point down (standard HNSW insertion).
        let mut layer_entry_point = global_entry_point;

        let entry_layer = (global_entry_layer as usize).min(max_layer);

        // 1) Descend from the current entry layer down to (current_layer + 1) using ef=1 greedy search.
        // We do NOT add edges here; we only refine the entry point.
        if entry_layer > current_layer {
            for layer_idx in ((current_layer + 1)..=entry_layer).rev() {
                let layer = &index.layers[layer_idx];
                let results = greedy_search_layer(
                    current_vector,
                    layer_entry_point,
                    layer,
                    &index.vectors,
                    index.dimension,
                    1,
                    dist_fn,
                );
                if let Some((best_id, _)) = results.first() {
                    layer_entry_point = *best_id;
                }
            }
        }

        // 2) For layers <= min(current_layer, entry_layer), run ef_construction search and connect.
        for layer_idx in (0..=current_layer.min(entry_layer)).rev() {
            let candidates = greedy_search_layer(
                current_vector,
                layer_entry_point,
                &index.layers[layer_idx],
                &index.vectors,
                index.dimension,
                index.params.ef_construction,
                dist_fn,
            );

            // Update entry point for the next lower layer.
            if let Some((best_id, _)) = candidates.first() {
                layer_entry_point = *best_id;
            }

            // Select neighbors using configured diversification strategy
            let m_actual = if layer_idx == 0 {
                index.params.m_max
            } else {
                index.params.m
            };

            let mut candidates = candidates;
            let mut selected = select_neighbors(
                current_vector,
                &mut candidates,
                m_actual,
                &index.vectors,
                index.dimension,
                &index.params.neighborhood_diversification,
                dist_fn,
            );

            // Enforce layer membership + insertion order invariants.
            // We are inserting `current_id` now, so only nodes < current_id exist in the graph.
            selected.retain(|&id| {
                let id_usize = id as usize;
                id_usize < current_id && (index.layer_assignments[id_usize] as usize) >= layer_idx
            });

            // Pre-compute neighbor distances (before any mutable borrows).
            // Vectors are looked up via get_vector(&index.vectors, ...) when needed,
            // avoiding a Vec<f32> heap allocation per neighbor.
            let neighbor_data: Vec<(u32, f32)> = selected
                .iter()
                .map(|&id| {
                    let vec = get_vector(&index.vectors, index.dimension, id as usize);
                    let dist = dist_fn(current_vector, vec);
                    (id, dist)
                })
                .collect();

            // Pre-compute existing neighbors of current_id
            let current_existing_neighbors: Vec<u32> = if layer_idx < index.layers.len() {
                index.layers[layer_idx]
                    .get_neighbors(current_id as u32)
                    .iter()
                    .copied()
                    .collect()
            } else {
                Vec::new()
            };

            // Pre-compute distances from current to ALL its neighbors (existing + selected).
            // Uses a flat Vec instead of HashMap: for m ≤ 32 entries, linear scan avoids
            // hashing overhead and repeated allocations that dominated 7-8% of construction.
            let mut all_current_distances: Vec<(u32, f32)> =
                Vec::with_capacity(current_existing_neighbors.len() + selected.len());

            // Add distances to existing neighbors
            for &id in &current_existing_neighbors {
                let vec = get_vector(&index.vectors, index.dimension, id as usize);
                let dist = dist_fn(current_vector, vec);
                all_current_distances.push((id, dist));
            }

            // Add distances to selected neighbors
            for &(nid, dist) in &neighbor_data {
                all_current_distances.push((nid, dist));
            }

            // Pre-compute distances for each selected neighbor to ALL its potential neighbors.
            // Inlines neighbor list iteration to avoid intermediate Vec<Vec<u32>> allocation.
            let mut all_reverse_distances: Vec<Vec<(u32, f32)>> =
                Vec::with_capacity(neighbor_data.len());
            for &(neighbor_id, dist_to_current) in &neighbor_data {
                let neighbor_vec =
                    get_vector(&index.vectors, index.dimension, neighbor_id as usize);
                let existing: Vec<u32> = if layer_idx < index.layers.len() {
                    index.layers[layer_idx]
                        .get_neighbors(neighbor_id)
                        .iter()
                        .copied()
                        .collect()
                } else {
                    Vec::new()
                };
                let mut distances = Vec::with_capacity(1 + existing.len());

                // Distance to current_id (already computed)
                distances.push((current_id as u32, dist_to_current));

                // Distances to existing neighbors
                for &existing_id in &existing {
                    let existing_vec =
                        get_vector(&index.vectors, index.dimension, existing_id as usize);
                    let dist = dist_fn(neighbor_vec, existing_vec);
                    distances.push((existing_id, dist));
                }

                all_reverse_distances.push(distances);
            }

            // Now do all mutable operations
            let layer = &mut index.layers[layer_idx];
            let neighbors_vec = layer.get_neighbors_mut();

            // First pass: add all edges without pruning
            for &neighbor_id in &selected {
                let neighbors = &mut neighbors_vec[current_id];
                if !neighbors.contains(&neighbor_id) {
                    neighbors.push(neighbor_id);
                }

                let reverse_neighbors = &mut neighbors_vec[neighbor_id as usize];
                if !reverse_neighbors.contains(&(current_id as u32)) {
                    reverse_neighbors.push(current_id as u32);
                }
            }

            // Second pass: prune current_id's neighbors using diversity heuristic
            // Paper (Algorithm 4): pruning should use the same heuristic as selection
            {
                let neighbors = &mut neighbors_vec[current_id];
                if neighbors.len() > m_actual {
                    let mut neighbor_candidates: Vec<(u32, f32)> = neighbors
                        .iter()
                        .map(|&id| {
                            let dist = all_current_distances
                                .iter()
                                .find(|(k, _)| *k == id)
                                .map(|(_, v)| *v)
                                .unwrap_or_else(|| {
                                    let vec =
                                        get_vector(&index.vectors, index.dimension, id as usize);
                                    dist_fn(current_vector, vec)
                                });
                            (id, dist)
                        })
                        .collect();

                    let pruned = select_neighbors(
                        current_vector,
                        &mut neighbor_candidates,
                        m_actual,
                        &index.vectors,
                        index.dimension,
                        &index.params.neighborhood_diversification,
                        dist_fn,
                    );
                    *neighbors = pruned.into_iter().collect();
                }
            }

            // Third pass: prune each selected neighbor's reverse list using diversity heuristic
            for (idx, &neighbor_id) in selected.iter().enumerate() {
                let reverse_neighbors = &mut neighbors_vec[neighbor_id as usize];
                if reverse_neighbors.len() > m_actual {
                    let distances = &all_reverse_distances[idx];
                    let neighbor_vec =
                        get_vector(&index.vectors, index.dimension, neighbor_id as usize);

                    let mut reverse_candidates: Vec<(u32, f32)> = reverse_neighbors
                        .iter()
                        .map(|&id| {
                            let dist = distances
                                .iter()
                                .find(|(k, _)| *k == id)
                                .map(|(_, v)| *v)
                                .unwrap_or_else(|| {
                                    let vec =
                                        get_vector(&index.vectors, index.dimension, id as usize);
                                    dist_fn(neighbor_vec, vec)
                                });
                            (id, dist)
                        })
                        .collect();

                    let pruned = select_neighbors(
                        neighbor_vec,
                        &mut reverse_candidates,
                        m_actual,
                        &index.vectors,
                        index.dimension,
                        &index.params.neighborhood_diversification,
                        dist_fn,
                    );
                    *reverse_neighbors = pruned.into_iter().collect();
                }
            }
        }

        // 3) Update global entry point if this node reaches a new top layer.
        if current_layer > (global_entry_layer as usize) {
            global_entry_point = current_id as u32;
            global_entry_layer = index.layer_assignments[current_id];
        }
    }

    Ok(())
}