engine/

advanced_search.rs

//! Advanced Search Features for Dakera
//!
//! Provides enhanced search capabilities:
//! - Multi-vector queries (query by multiple vectors)
//! - Negative vectors (exclusion/avoidance)
//! - MMR (Maximal Marginal Relevance) for diversity
//! - Range queries (distance threshold)
//! - Result aggregation and grouping

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

use common::{DistanceMetric, VectorId};

/// Configuration for advanced search operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdvancedSearchConfig {
    /// Enable MMR diversity ranking
    pub enable_mmr: bool,
    /// Lambda parameter for MMR (0 = max diversity, 1 = max relevance)
    pub mmr_lambda: f32,
    /// Maximum candidates to consider for MMR reranking
    pub mmr_candidates: usize,
    /// Enable result grouping
    pub enable_grouping: bool,
    /// Field to group results by
    pub group_by_field: Option<String>,
    /// Maximum results per group
    pub max_per_group: usize,
}

impl Default for AdvancedSearchConfig {
    fn default() -> Self {
        Self {
            enable_mmr: false,
            mmr_lambda: 0.5,
            mmr_candidates: 100,
            enable_grouping: false,
            group_by_field: None,
            max_per_group: 3,
        }
    }
}

/// Multi-vector query specification
#[derive(Debug, Clone)]
pub struct MultiVectorQuery {
    /// Positive vectors (search towards these)
    pub positive_vectors: Vec<Vec<f32>>,
    /// Weights for positive vectors
    pub positive_weights: Vec<f32>,
    /// Negative vectors (search away from these)
    pub negative_vectors: Vec<Vec<f32>>,
    /// Weights for negative vectors
    pub negative_weights: Vec<f32>,
    /// Number of results to return
    pub top_k: usize,
    /// Optional distance threshold
    pub distance_threshold: Option<f32>,
}

impl MultiVectorQuery {
    /// Create a simple single-vector query
    pub fn single(vector: Vec<f32>, top_k: usize) -> Self {
        Self {
            positive_vectors: vec![vector],
            positive_weights: vec![1.0],
            negative_vectors: Vec::new(),
            negative_weights: Vec::new(),
            top_k,
            distance_threshold: None,
        }
    }

    /// Create a multi-vector query
    pub fn multi(vectors: Vec<Vec<f32>>, top_k: usize) -> Self {
        let weights = vec![1.0 / vectors.len() as f32; vectors.len()];
        Self {
            positive_vectors: vectors,
            positive_weights: weights,
            negative_vectors: Vec::new(),
            negative_weights: Vec::new(),
            top_k,
            distance_threshold: None,
        }
    }

    /// Add a negative (avoidance) vector
    pub fn with_negative(mut self, vector: Vec<f32>, weight: f32) -> Self {
        self.negative_vectors.push(vector);
        self.negative_weights.push(weight);
        self
    }

    /// Set distance threshold for range query
    pub fn with_threshold(mut self, threshold: f32) -> Self {
        self.distance_threshold = Some(threshold);
        self
    }

    /// Set custom weights for positive vectors
    pub fn with_weights(mut self, weights: Vec<f32>) -> Self {
        self.positive_weights = weights;
        self
    }

    /// Compute the effective query vector (weighted combination)
    pub fn compute_query_vector(&self, dimensions: usize) -> Vec<f32> {
        let mut result = vec![0.0; dimensions];

        // Add weighted positive vectors
        for (vec, &weight) in self.positive_vectors.iter().zip(&self.positive_weights) {
            for (i, &v) in vec.iter().enumerate() {
                if i < dimensions {
                    result[i] += v * weight;
                }
            }
        }

        // Subtract weighted negative vectors
        for (vec, &weight) in self.negative_vectors.iter().zip(&self.negative_weights) {
            for (i, &v) in vec.iter().enumerate() {
                if i < dimensions {
                    result[i] -= v * weight;
                }
            }
        }

        // Normalize the result
        let norm: f32 = result.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for v in &mut result {
                *v /= norm;
            }
        }

        result
    }
}

/// Result from advanced search
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdvancedSearchResult {
    /// Vector ID
    pub id: VectorId,
    /// Similarity/distance score
    pub score: f32,
    /// Original rank before reranking
    pub original_rank: usize,
    /// Final rank after reranking
    pub final_rank: usize,
    /// MMR score (if enabled)
    pub mmr_score: Option<f32>,
    /// Group key (if grouping enabled)
    pub group_key: Option<String>,
}

/// MMR (Maximal Marginal Relevance) reranker for diversity
pub struct MmrReranker {
    /// Lambda parameter (0 = max diversity, 1 = max relevance)
    lambda: f32,
}

impl MmrReranker {
    /// Create a new MMR reranker
    pub fn new(lambda: f32) -> Self {
        Self {
            lambda: lambda.clamp(0.0, 1.0),
        }
    }

    /// Rerank results using MMR for diversity
    ///
    /// MMR = λ * Sim(d, q) - (1 - λ) * max(Sim(d, d_j))
    /// where d_j are already selected documents
    pub fn rerank(
        &self,
        candidates: &[(VectorId, f32, Vec<f32>)], // (id, score, vector)
        top_k: usize,
    ) -> Vec<AdvancedSearchResult> {
        if candidates.is_empty() {
            return Vec::new();
        }

        let mut selected: Vec<usize> = Vec::with_capacity(top_k);
        let mut remaining: HashSet<usize> = (0..candidates.len()).collect();
        let mut results = Vec::with_capacity(top_k);

        // Select first item (highest relevance)
        let first_idx = candidates
            .iter()
            .enumerate()
            .max_by(|a, b| {
                a.1 .1
                    .partial_cmp(&b.1 .1)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|(i, _)| i)
            .unwrap_or(0);

        selected.push(first_idx);
        remaining.remove(&first_idx);
        results.push(AdvancedSearchResult {
            id: candidates[first_idx].0.clone(),
            score: candidates[first_idx].1,
            original_rank: first_idx,
            final_rank: 0,
            mmr_score: Some(candidates[first_idx].1),
            group_key: None,
        });

        // Iteratively select remaining items
        while results.len() < top_k && !remaining.is_empty() {
            let mut best_idx = None;
            let mut best_mmr = f32::NEG_INFINITY;

            for &idx in &remaining {
                let relevance = candidates[idx].1;

                // Compute max similarity to already selected items
                let max_sim = selected
                    .iter()
                    .map(|&sel_idx| {
                        self.cosine_similarity(&candidates[idx].2, &candidates[sel_idx].2)
                    })
                    .fold(f32::NEG_INFINITY, f32::max);

                // MMR score
                let mmr = self.lambda * relevance - (1.0 - self.lambda) * max_sim;

                if mmr > best_mmr {
                    best_mmr = mmr;
                    best_idx = Some(idx);
                }
            }

            if let Some(idx) = best_idx {
                selected.push(idx);
                remaining.remove(&idx);
                results.push(AdvancedSearchResult {
                    id: candidates[idx].0.clone(),
                    score: candidates[idx].1,
                    original_rank: idx,
                    final_rank: results.len(),
                    mmr_score: Some(best_mmr),
                    group_key: None,
                });
            } else {
                break;
            }
        }

        results
    }

    /// Compute cosine similarity between two vectors
    fn cosine_similarity(&self, a: &[f32], b: &[f32]) -> f32 {
        let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
        let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

        if norm_a > 0.0 && norm_b > 0.0 {
            dot / (norm_a * norm_b)
        } else {
            0.0
        }
    }
}

/// Range query executor for distance-based filtering
pub struct RangeQuery {
    /// Distance metric to use
    metric: DistanceMetric,
    /// Maximum distance threshold
    threshold: f32,
}

impl RangeQuery {
    /// Create a new range query
    pub fn new(metric: DistanceMetric, threshold: f32) -> Self {
        Self { metric, threshold }
    }

    /// Filter results by distance threshold
    pub fn filter(&self, results: Vec<(VectorId, f32)>) -> Vec<(VectorId, f32)> {
        results
            .into_iter()
            .filter(|(_, score)| self.passes_threshold(*score))
            .collect()
    }

    /// Check if a score passes the threshold
    fn passes_threshold(&self, score: f32) -> bool {
        match self.metric {
            // For cosine/dot product, higher is better
            DistanceMetric::Cosine | DistanceMetric::DotProduct => score >= self.threshold,
            // For euclidean, lower is better (negative euclidean, so higher is better)
            DistanceMetric::Euclidean => score >= -self.threshold,
        }
    }
}

/// Result grouper for organizing results by field
pub struct ResultGrouper {
    /// Field to group by
    group_field: String,
    /// Maximum results per group
    max_per_group: usize,
}

impl ResultGrouper {
    /// Create a new result grouper
    pub fn new(group_field: String, max_per_group: usize) -> Self {
        Self {
            group_field,
            max_per_group,
        }
    }

    /// Group results by field value
    pub fn group(
        &self,
        results: Vec<(VectorId, f32, Option<serde_json::Value>)>,
    ) -> HashMap<String, Vec<(VectorId, f32)>> {
        let mut groups: HashMap<String, Vec<(VectorId, f32)>> = HashMap::new();

        for (id, score, metadata) in results {
            let group_key = metadata
                .and_then(|m| m.get(&self.group_field).cloned())
                .and_then(|v| match v {
                    serde_json::Value::String(s) => Some(s),
                    serde_json::Value::Number(n) => Some(n.to_string()),
                    _ => None,
                })
                .unwrap_or_else(|| "_ungrouped".to_string());

            let group = groups.entry(group_key).or_default();
            if group.len() < self.max_per_group {
                group.push((id, score));
            }
        }

        groups
    }
}

/// Advanced search executor combining all features
pub struct AdvancedSearchExecutor {
    config: AdvancedSearchConfig,
}

impl AdvancedSearchExecutor {
    /// Create a new advanced search executor
    pub fn new(config: AdvancedSearchConfig) -> Self {
        Self { config }
    }

    /// Process search results with advanced features
    pub fn process_results(
        &self,
        candidates: Vec<(VectorId, f32, Vec<f32>)>,
        query: &MultiVectorQuery,
    ) -> Vec<AdvancedSearchResult> {
        let mut results = candidates;

        // Apply distance threshold if specified
        if let Some(threshold) = query.distance_threshold {
            results.retain(|(_, score, _)| *score >= threshold);
        }

        // Apply MMR if enabled
        if self.config.enable_mmr {
            let reranker = MmrReranker::new(self.config.mmr_lambda);
            return reranker.rerank(&results, query.top_k);
        }

        // Otherwise, return top-k by score
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        results.truncate(query.top_k);

        results
            .into_iter()
            .enumerate()
            .map(|(rank, (id, score, _))| AdvancedSearchResult {
                id,
                score,
                original_rank: rank,
                final_rank: rank,
                mmr_score: None,
                group_key: None,
            })
            .collect()
    }

    /// Apply negative vector penalty to scores
    pub fn apply_negative_penalty(
        &self,
        results: &mut [(VectorId, f32, Vec<f32>)],
        negative_vectors: &[Vec<f32>],
        negative_weights: &[f32],
    ) {
        for (_, score, vec) in results.iter_mut() {
            for (neg_vec, &weight) in negative_vectors.iter().zip(negative_weights) {
                // Compute similarity to negative vector
                let sim = self.cosine_similarity(vec, neg_vec);
                // Penalize score based on similarity to negative vector
                *score -= sim * weight;
            }
        }
    }

    fn cosine_similarity(&self, a: &[f32], b: &[f32]) -> f32 {
        let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
        let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

        if norm_a > 0.0 && norm_b > 0.0 {
            dot / (norm_a * norm_b)
        } else {
            0.0
        }
    }
}

/// Search statistics for monitoring
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SearchStats {
    /// Total candidates considered
    pub candidates_considered: usize,
    /// Results after threshold filtering
    pub after_threshold: usize,
    /// Results after MMR reranking
    pub after_mmr: usize,
    /// Number of groups (if grouping enabled)
    pub num_groups: usize,
    /// Search latency in milliseconds
    pub latency_ms: u64,
}

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

    #[test]
    fn test_multi_vector_query() {
        let query = MultiVectorQuery::multi(vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]], 10);

        assert_eq!(query.positive_vectors.len(), 2);
        assert_eq!(query.positive_weights.len(), 2);
        assert_eq!(query.positive_weights[0], 0.5);
    }

    #[test]
    fn test_query_vector_computation() {
        let query = MultiVectorQuery::single(vec![1.0, 0.0, 0.0], 10);
        let computed = query.compute_query_vector(3);

        assert_eq!(computed.len(), 3);
        assert!((computed[0] - 1.0).abs() < 0.01);
    }

    #[test]
    fn test_negative_vector() {
        let query = MultiVectorQuery::single(vec![1.0, 0.0, 0.0], 10)
            .with_negative(vec![0.0, 1.0, 0.0], 0.5);

        assert_eq!(query.negative_vectors.len(), 1);
        assert_eq!(query.negative_weights[0], 0.5);
    }

    #[test]
    fn test_mmr_reranker() {
        let reranker = MmrReranker::new(0.5);

        // Create candidates with varying similarities
        let candidates = vec![
            ("a".to_string(), 0.9, vec![1.0, 0.0, 0.0]),
            ("b".to_string(), 0.85, vec![0.95, 0.1, 0.0]), // Similar to a
            ("c".to_string(), 0.8, vec![0.0, 1.0, 0.0]),   // Different from a
            ("d".to_string(), 0.75, vec![0.0, 0.0, 1.0]),  // Different from both
        ];

        let results = reranker.rerank(&candidates, 3);

        assert_eq!(results.len(), 3);
        // First should be highest relevance
        assert_eq!(results[0].id, "a");
        // Due to MMR, "c" or "d" should rank higher than "b" (more diverse)
    }

    #[test]
    fn test_range_query() {
        let range = RangeQuery::new(DistanceMetric::Cosine, 0.8);

        let results = vec![
            ("a".to_string(), 0.95),
            ("b".to_string(), 0.75), // Below threshold
            ("c".to_string(), 0.85),
        ];

        let filtered = range.filter(results);
        assert_eq!(filtered.len(), 2);
        assert!(filtered.iter().all(|(_, s)| *s >= 0.8));
    }

    #[test]
    fn test_result_grouper() {
        let grouper = ResultGrouper::new("category".to_string(), 2);

        let results = vec![
            (
                "a".to_string(),
                0.9,
                Some(serde_json::json!({"category": "tech"})),
            ),
            (
                "b".to_string(),
                0.85,
                Some(serde_json::json!({"category": "tech"})),
            ),
            (
                "c".to_string(),
                0.8,
                Some(serde_json::json!({"category": "tech"})),
            ), // Should be excluded (max 2)
            (
                "d".to_string(),
                0.75,
                Some(serde_json::json!({"category": "science"})),
            ),
        ];

        let groups = grouper.group(results);

        assert_eq!(groups.len(), 2);
        assert_eq!(groups["tech"].len(), 2);
        assert_eq!(groups["science"].len(), 1);
    }

    #[test]
    fn test_advanced_search_executor() {
        let config = AdvancedSearchConfig {
            enable_mmr: false,
            ..Default::default()
        };
        let executor = AdvancedSearchExecutor::new(config);

        let candidates = vec![
            ("a".to_string(), 0.9, vec![1.0, 0.0]),
            ("b".to_string(), 0.8, vec![0.0, 1.0]),
            ("c".to_string(), 0.7, vec![0.5, 0.5]),
        ];

        let query = MultiVectorQuery::single(vec![1.0, 0.0], 2);
        let results = executor.process_results(candidates, &query);

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, "a");
        assert_eq!(results[1].id, "b");
    }

    #[test]
    fn test_threshold_filtering() {
        let config = AdvancedSearchConfig::default();
        let executor = AdvancedSearchExecutor::new(config);

        let candidates = vec![
            ("a".to_string(), 0.9, vec![1.0, 0.0]),
            ("b".to_string(), 0.5, vec![0.0, 1.0]), // Below threshold
            ("c".to_string(), 0.85, vec![0.5, 0.5]),
        ];

        let query = MultiVectorQuery::single(vec![1.0, 0.0], 10).with_threshold(0.7);
        let results = executor.process_results(candidates, &query);

        assert_eq!(results.len(), 2);
        assert!(results.iter().all(|r| r.score >= 0.7));
    }
}