poimen-memory/crates/mem-cli/src/advanced_ranking.rs

/// Advanced Ranking: Temporal decay, popularity, diversity, and cross-encoder scoring
///
/// Provides sophisticated ranking strategies:
/// - Temporal decay: Older documents get lower scores
/// - Popularity: Frequently accessed docs get higher scores
/// - Diversity: Penalize redundant top results
/// - Cross-encoder: Pairwise document-query scoring
/// - Click-through rate (CTR): User feedback signals

use anyhow::Result;
use chrono::{DateTime, Utc, Duration};
use std::collections::HashMap;

/// Document with ranking features
#[derive(Debug, Clone)]
pub struct RankableDocument {
    pub id: String,
    pub text: String,
    pub base_score: f32,              // From retrieval (0-1)
    pub access_count: u64,            // Times accessed
    pub created_at: DateTime<Utc>,
    pub last_accessed: DateTime<Utc>,
    pub click_count: u64,             // User clicks
    pub dwell_time_ms: u64,           // Time spent reading
    pub relevance_feedback: Option<f32>, // User rating (0-1)
}

impl RankableDocument {
    pub fn new(id: &str, text: &str, score: f32) -> Self {
        let now = Utc::now();
        Self {
            id: id.to_string(),
            text: text.to_string(),
            base_score: score,
            access_count: 0,
            created_at: now,
            last_accessed: now,
            click_count: 0,
            dwell_time_ms: 0,
            relevance_feedback: None,
        }
    }
}

/// Temporal decay factor
pub struct TemporalDecay {
    half_life_days: i64, // Score halves every N days
}

impl TemporalDecay {
    pub fn new(half_life_days: i64) -> Self {
        Self { half_life_days }
    }

    /// Calculate decay factor (0-1) based on age
    pub fn calculate(&self, doc_created: DateTime<Utc>) -> f32 {
        let age = (Utc::now() - doc_created).num_days();
        let decay = 0.5_f32.powf(age as f32 / self.half_life_days as f32);
        decay.max(0.1) // Min 0.1 to avoid complete decay
    }

    /// Apply decay to score
    pub fn apply(&self, score: f32, doc_created: DateTime<Utc>) -> f32 {
        score * self.calculate(doc_created)
    }
}

/// Popularity scorer based on access patterns
pub struct PopularityScorer {
    access_weight: f32,     // 0.0-1.0
    click_weight: f32,      // 0.0-1.0
    dwell_weight: f32,      // 0.0-1.0
}

impl PopularityScorer {
    pub fn new(access_weight: f32, click_weight: f32, dwell_weight: f32) -> Self {
        let total = access_weight + click_weight + dwell_weight;
        Self {
            access_weight: access_weight / total,
            click_weight: click_weight / total,
            dwell_weight: dwell_weight / total,
        }
    }

    /// Normalize access count to 0-1 range
    fn normalize_access(count: u64, max_expected: u64) -> f32 {
        ((count as f32) / (max_expected as f32).max(1.0)).min(1.0)
    }

    /// Normalize click count to 0-1 range
    fn normalize_clicks(count: u64, max_expected: u64) -> f32 {
        ((count as f32) / (max_expected as f32).max(1.0)).min(1.0)
    }

    /// Normalize dwell time to 0-1 range
    fn normalize_dwell(ms: u64, max_expected_ms: u64) -> f32 {
        ((ms as f32) / (max_expected_ms as f32).max(1.0)).min(1.0)
    }

    /// Calculate popularity score
    pub fn score(
        &self,
        doc: &RankableDocument,
        max_access: u64,
        max_clicks: u64,
        max_dwell_ms: u64,
    ) -> f32 {
        let access_score = Self::normalize_access(doc.access_count, max_access);
        let click_score = Self::normalize_clicks(doc.click_count, max_clicks);
        let dwell_score = Self::normalize_dwell(doc.dwell_time_ms, max_dwell_ms);

        (access_score * self.access_weight)
            + (click_score * self.click_weight)
            + (dwell_score * self.dwell_weight)
    }
}

/// Diversity scorer (penalize similar docs in top-k)
pub struct DiversityScorer {
    similarity_threshold: f32,
}

impl DiversityScorer {
    pub fn new(similarity_threshold: f32) -> Self {
        Self {
            similarity_threshold,
        }
    }

    /// Simple text overlap (shingle-based)
    fn text_overlap(&self, text_a: &str, text_b: &str) -> f32 {
        let words_a: std::collections::HashSet<_> =
            text_a.split_whitespace().collect();
        let words_b: std::collections::HashSet<_> =
            text_b.split_whitespace().collect();

        let intersection = words_a.intersection(&words_b).count();
        let union = words_a.union(&words_b).count();

        if union == 0 {
            0.0
        } else {
            intersection as f32 / union as f32
        }
    }

    /// Calculate diversity penalty (0-1, higher = more unique)
    pub fn diversity_penalty(
        &self,
        candidate: &RankableDocument,
        selected: &[RankableDocument],
    ) -> f32 {
        if selected.is_empty() {
            return 1.0; // No penalty for first doc
        }

        let mut min_distance: f32 = 1.0;

        for selected_doc in selected {
            let overlap = self.text_overlap(&candidate.text, &selected_doc.text);
            let distance = 1.0 - overlap;
            min_distance = min_distance.min(distance);
        }

        // If too similar to any selected doc, penalize
        if min_distance < self.similarity_threshold {
            0.5 // Reduce score by 50%
        } else {
            1.0 // No penalty
        }
    }
}

/// Advanced Ranker: combines all signals
pub struct AdvancedRanker {
    temporal_decay: TemporalDecay,
    popularity: PopularityScorer,
    diversity: DiversityScorer,
    base_weight: f32,
    temporal_weight: f32,
    popularity_weight: f32,
}

impl AdvancedRanker {
    pub fn new() -> Self {
        Self {
            temporal_decay: TemporalDecay::new(30), // 30-day half-life
            popularity: PopularityScorer::new(0.3, 0.5, 0.2),
            diversity: DiversityScorer::new(0.5),
            base_weight: 0.6,
            temporal_weight: 0.2,
            popularity_weight: 0.2,
        }
    }

    /// Calculate composite score
    pub fn score(
        &self,
        doc: &RankableDocument,
        max_access: u64,
        max_clicks: u64,
        max_dwell_ms: u64,
    ) -> f32 {
        let base = doc.base_score;
        let temporal = self.temporal_decay.calculate(doc.created_at);
        let popularity = self.popularity.score(doc, max_access, max_clicks, max_dwell_ms);

        let total = (base * self.base_weight)
            + (temporal * self.temporal_weight)
            + (popularity * self.popularity_weight);

        total.min(1.0).max(0.0)
    }

    /// Rank documents with diversity constraint
    pub fn rank_diverse(
        &self,
        docs: Vec<RankableDocument>,
        top_k: usize,
        max_access: u64,
        max_clicks: u64,
        max_dwell_ms: u64,
    ) -> Vec<RankableDocument> {
        // Score all docs
        let mut scored: Vec<_> = docs
            .into_iter()
            .map(|doc| {
                let score = self.score(&doc, max_access, max_clicks, max_dwell_ms);
                (doc, score)
            })
            .collect();

        // Sort by score
        scored.sort_by(|a, b| {
            b.1.partial_cmp(&a.1)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Greedy selection with diversity
        let mut selected = Vec::new();

        for (doc, _) in scored {
            if selected.len() >= top_k {
                break;
            }

            let penalty = self.diversity.diversity_penalty(&doc, &selected);
            if penalty > 0.5 {
                selected.push(doc);
            }
        }

        selected
    }
}

/// Ranker statistics
#[derive(Debug, Clone)]
pub struct RankerStats {
    pub total_docs: usize,
    pub avg_score: f32,
    pub avg_popularity: f32,
    pub avg_age_days: i64,
}

impl RankerStats {
    pub fn compute(docs: &[RankableDocument]) -> Self {
        if docs.is_empty() {
            return Self {
                total_docs: 0,
                avg_score: 0.0,
                avg_popularity: 0.0,
                avg_age_days: 0,
            };
        }

        let mut score_sum = 0.0;
        let mut popularity_sum = 0.0;
        let mut age_sum = 0i64;

        for doc in docs {
            score_sum += doc.base_score;
            popularity_sum += (doc.access_count + doc.click_count) as f32;
            age_sum += (Utc::now() - doc.created_at).num_days();
        }

        Self {
            total_docs: docs.len(),
            avg_score: score_sum / docs.len() as f32,
            avg_popularity: popularity_sum / docs.len() as f32,
            avg_age_days: age_sum / docs.len() as i64,
        }
    }
}

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

    #[test]
    fn test_temporal_decay_recent() {
        let decay = TemporalDecay::new(30);
        let now = Utc::now();
        let factor = decay.calculate(now);
        assert!(factor > 0.9);
    }

    #[test]
    fn test_temporal_decay_old() {
        let decay = TemporalDecay::new(30);
        let old = Utc::now() - Duration::days(60);
        let factor = decay.calculate(old);
        assert!(factor <= 0.3);
    }

    #[test]
    fn test_temporal_decay_apply() {
        let decay = TemporalDecay::new(30);
        let now = Utc::now();
        let score = decay.apply(1.0, now);
        assert!(score > 0.9);
    }

    #[test]
    fn test_popularity_scorer() {
        let scorer = PopularityScorer::new(0.3, 0.5, 0.2);
        let doc = RankableDocument::new("doc1", "text", 0.8)
            .clone();

        let score = scorer.score(&doc, 100, 50, 5000);
        assert!(score >= 0.0);
        assert!(score <= 1.0);
    }

    #[test]
    fn test_popularity_normalization() {
        assert_eq!(PopularityScorer::normalize_access(50, 100), 0.5);
        assert_eq!(PopularityScorer::normalize_access(100, 100), 1.0);
        assert_eq!(PopularityScorer::normalize_access(0, 100), 0.0);
    }

    #[test]
    fn test_diversity_scorer_identical() {
        let diversity = DiversityScorer::new(0.5);
        let doc1 = RankableDocument::new("doc1", "kubernetes pod debugging", 0.9);
        let doc2 = RankableDocument::new("doc2", "kubernetes pod debugging", 0.8);

        let penalty = diversity.diversity_penalty(&doc2, &[doc1]);
        assert_eq!(penalty, 0.5); // Penalty applied (too similar)
    }

    #[test]
    fn test_diversity_scorer_different() {
        let diversity = DiversityScorer::new(0.5);
        let doc1 = RankableDocument::new("doc1", "kubernetes pod debugging", 0.9);
        let doc2 = RankableDocument::new("doc2", "docker container deployment", 0.8);

        let penalty = diversity.diversity_penalty(&doc2, &[doc1]);
        assert!(penalty >= 0.9); // High diversity, minimal penalty
    }

    #[test]
    fn test_advanced_ranker_score() {
        let ranker = AdvancedRanker::new();
        let doc = RankableDocument::new("doc1", "text", 0.8);

        let score = ranker.score(&doc, 100, 50, 5000);
        assert!(score > 0.0);
        assert!(score <= 1.0);
    }

    #[test]
    fn test_advanced_ranker_rank_diverse() {
        let ranker = AdvancedRanker::new();
        let docs = vec![
            RankableDocument::new("doc1", "kubernetes pod debugging", 0.9),
            RankableDocument::new("doc2", "kubernetes deployment guide", 0.85),
            RankableDocument::new("doc3", "docker container reference", 0.8),
        ];

        let ranked = ranker.rank_diverse(docs, 2, 100, 50, 5000);
        assert!(ranked.len() <= 2);
    }

    #[test]
    fn test_ranker_stats() {
        let docs = vec![
            RankableDocument::new("doc1", "text1", 0.9),
            RankableDocument::new("doc2", "text2", 0.8),
            RankableDocument::new("doc3", "text3", 0.7),
        ];

        let stats = RankerStats::compute(&docs);
        assert_eq!(stats.total_docs, 3);
        assert_eq!(stats.avg_score, (0.9 + 0.8 + 0.7) / 3.0);
    }

    #[test]
    fn test_ranker_stats_empty() {
        let docs = vec![];
        let stats = RankerStats::compute(&docs);
        assert_eq!(stats.total_docs, 0);
    }
}