motedb 0.7.8

//! Robust Prune algorithm for Vamana
//!
//! The pruning strategy maintains graph connectivity while limiting node degree.

use crate::types::RowId;
use std::cmp::Ordering;

/// Candidate neighbor with distance
#[derive(Debug, Clone)]
pub struct Candidate {
    pub id: RowId,
    pub distance: f32,
}

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

impl Eq for Candidate {}

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

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

/// Robust Prune algorithm
///
/// Prunes a set of candidates to at most R neighbors while maintaining diversity.
///
/// # Arguments
/// * `candidates` - List of candidate neighbors with distances
/// * `max_degree` - Maximum number of neighbors (R parameter)
/// * `alpha` - Diversity parameter (typically 1.2)
/// * `distance_fn` - Function to compute distance between two vectors
///
/// # Returns
/// Pruned list of at most max_degree neighbors
pub fn robust_prune<F>(
    candidates: Vec<Candidate>,
    max_degree: usize,
    alpha: f32,
    distance_fn: F,
) -> Vec<RowId>
where
    F: Fn(RowId, RowId) -> f32,
{
    if candidates.len() <= max_degree {
        return candidates.into_iter().map(|c| c.id).collect();
    }

    // Sort candidates by distance (ascending)
    let mut sorted_candidates: Vec<_> = candidates.into_iter().collect();
    sorted_candidates.sort_by(|a, b| {
        a.distance
            .partial_cmp(&b.distance)
            .unwrap_or(Ordering::Equal)
    });

    let mut pruned = Vec::with_capacity(max_degree);

    for candidate in sorted_candidates {
        if pruned.len() >= max_degree {
            break;
        }

        // Check diversity constraint
        let mut should_add = true;

        for &selected_id in &pruned {
            let dist_to_selected = distance_fn(candidate.id, selected_id);

            // If candidate is too close to an already selected neighbor,
            // and the selected neighbor is closer to the target, skip this candidate
            if dist_to_selected < alpha * candidate.distance {
                should_add = false;
                break;
            }
        }

        if should_add {
            pruned.push(candidate.id);
        }
    }

    pruned
}

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

    #[test]
    fn test_candidate_ordering() {
        let c1 = Candidate {
            id: 1,
            distance: 1.0,
        };
        let c2 = Candidate {
            id: 2,
            distance: 2.0,
        };

        // c1 should be "greater" (min-heap)
        assert!(c1 > c2);
    }

    #[test]
    fn test_robust_prune_basic() {
        let candidates = vec![
            Candidate {
                id: 1,
                distance: 1.0,
            },
            Candidate {
                id: 2,
                distance: 2.0,
            },
            Candidate {
                id: 3,
                distance: 3.0,
            },
        ];

        // Mock distance function (returns constant)
        let dist_fn = |_a: RowId, _b: RowId| 10.0;

        let pruned = robust_prune(candidates, 2, 1.2, dist_fn);

        assert_eq!(pruned.len(), 2);
        assert_eq!(pruned[0], 1); // Closest should always be included
    }

    #[test]
    fn test_robust_prune_diversity() {
        let candidates = vec![
            Candidate {
                id: 1,
                distance: 1.0,
            },
            Candidate {
                id: 2,
                distance: 1.1,
            }, // Very close to id:1
            Candidate {
                id: 3,
                distance: 5.0,
            },
        ];

        // Distance function: id:1 and id:2 are very close
        let dist_fn = |a: RowId, b: RowId| {
            if (a == 1 && b == 2) || (a == 2 && b == 1) {
                0.5 // Close together
            } else {
                10.0 // Far apart
            }
        };

        let pruned = robust_prune(candidates, 2, 1.2, dist_fn);

        // With alpha=1.2, id:2 should be pruned because:
        // dist(2, 1) = 0.5 < 1.2 * 1.1 = 1.32
        assert!(pruned.contains(&1));
        assert!(pruned.contains(&3));
    }
}