// Reciprocal Rank Fusion (RRF) Algorithm
// Combines multiple ranked lists into a single fused ranking
// Based on: "Reciprocal Rank Fusion outperforms Condorcet and individual Rank
// Learning Methods"

use std::collections::HashMap;

use tracing::debug;

/// RRF-fused result with combined score
#[derive(Debug, Clone)]
pub struct RRFResult {
    pub chunk_id: String,
    pub score: f32,
}

/// Reciprocal Rank Fusion configuration
#[derive(Debug, Clone)]
pub struct RRFConfig {
    /// K parameter for RRF formula (default: 60)
    pub k: f32,
}

impl Default for RRFConfig {
    fn default() -> Self {
        Self { k: 60.0 }
    }
}

/// Reciprocal Rank Fusion
///
/// Combines multiple ranked lists using the formula:
/// RRF(d) = sum_{r in R} 1 / (k + rank_r(d))
///
/// Where:
/// - d is a document (chunk)
/// - R is the set of ranking functions
/// - rank_r(d) is the rank of document d in ranking r
/// - k is a constant (typically 60)
///
/// # Example
///
/// ```rust,ignore
/// let bm25_results = vec![("chunk-1", 1), ("chunk-2", 2)];
/// let semantic_results = vec![("chunk-2", 1), ("chunk-1", 2)];
///
/// let fused = reciprocal_rank_fusion(
///     &[bm25_results, semantic_results],
///     RRFConfig::default(),
///     10,
/// );
/// ```
pub fn reciprocal_rank_fusion(
    ranked_lists: &[Vec<String>],
    config: RRFConfig,
    limit: usize,
) -> Vec<RRFResult> {
    if ranked_lists.is_empty() {
        return Vec::new();
    }

    let mut scores: HashMap<String, f32> = HashMap::new();

    // Compute RRF scores
    for ranked_list in ranked_lists {
        for (rank, chunk_id) in ranked_list.iter().enumerate() {
            let rrf_contribution = 1.0 / (config.k + (rank + 1) as f32);
            *scores.entry(chunk_id.clone()).or_insert(0.0) += rrf_contribution;
        }
    }

    // Convert to results and sort by score descending
    let mut results: Vec<RRFResult> = scores
        .into_iter()
        .map(|(chunk_id, score)| RRFResult { chunk_id, score })
        .collect();

    results.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // Take top N
    let top_results: Vec<RRFResult> = results.into_iter().take(limit).collect();

    debug!(
        "RRF fusion combined {} lists into {} results",
        ranked_lists.len(),
        top_results.len()
    );

    top_results
}

/// Reciprocal Rank Fusion with scores
///
/// Similar to `reciprocal_rank_fusion` but takes scored results
/// and uses the ranking (position in list) for RRF, not the raw scores.
pub fn reciprocal_rank_fusion_scored<T>(
    scored_lists: &[Vec<(String, T)>],
    config: RRFConfig,
    limit: usize,
) -> Vec<RRFResult> {
    if scored_lists.is_empty() {
        return Vec::new();
    }

    let mut scores: HashMap<String, f32> = HashMap::new();

    // Compute RRF scores (ignoring original scores, using only rank)
    for scored_list in scored_lists {
        for (rank, (chunk_id, _score)) in scored_list.iter().enumerate() {
            let rrf_contribution = 1.0 / (config.k + (rank + 1) as f32);
            *scores.entry(chunk_id.clone()).or_insert(0.0) += rrf_contribution;
        }
    }

    // Convert to results and sort by score descending
    let mut results: Vec<RRFResult> = scores
        .into_iter()
        .map(|(chunk_id, score)| RRFResult { chunk_id, score })
        .collect();

    results.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    // Take top N
    let top_results: Vec<RRFResult> = results.into_iter().take(limit).collect();

    debug!(
        "RRF fusion (scored) combined {} lists into {} results",
        scored_lists.len(),
        top_results.len()
    );

    top_results
}

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

    #[test]
    fn test_rrf_basic() {
        let list1 = vec![
            "chunk-1".to_string(),
            "chunk-2".to_string(),
            "chunk-3".to_string(),
        ];
        let list2 = vec![
            "chunk-2".to_string(),
            "chunk-1".to_string(),
            "chunk-4".to_string(),
        ];

        let results = reciprocal_rank_fusion(&[list1, list2], RRFConfig::default(), 10);

        assert!(!results.is_empty());
        // chunk-1 and chunk-2 appear in both lists, should rank higher
        assert!(results[0].chunk_id == "chunk-1" || results[0].chunk_id == "chunk-2");
    }

    #[test]
    fn test_rrf_empty_lists() {
        let results = reciprocal_rank_fusion(&[], RRFConfig::default(), 10);
        assert!(results.is_empty());
    }

    #[test]
    fn test_rrf_single_list() {
        let list = vec![
            "chunk-1".to_string(),
            "chunk-2".to_string(),
            "chunk-3".to_string(),
        ];

        let results = reciprocal_rank_fusion(&[list], RRFConfig::default(), 2);

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].chunk_id, "chunk-1"); // Highest rank
        assert_eq!(results[1].chunk_id, "chunk-2");
    }

    #[test]
    fn test_rrf_limit() {
        let list1 = vec![
            "chunk-1".to_string(),
            "chunk-2".to_string(),
            "chunk-3".to_string(),
            "chunk-4".to_string(),
        ];

        let results = reciprocal_rank_fusion(&[list1], RRFConfig::default(), 2);

        assert_eq!(results.len(), 2); // Limited to 2
    }

    #[test]
    fn test_rrf_score_calculation() {
        let list1 = vec!["chunk-1".to_string()];
        let list2 = vec!["chunk-1".to_string()];

        let config = RRFConfig { k: 60.0 };
        let results = reciprocal_rank_fusion(&[list1, list2], config, 10);

        assert_eq!(results.len(), 1);
        // RRF score for chunk-1: (1 / (60 + 1)) + (1 / (60 + 1)) = 2/61 ≈ 0.0328
        let expected_score = 2.0 / 61.0;
        assert!((results[0].score - expected_score).abs() < 0.001);
    }

    #[test]
    fn test_rrf_scored_variant() {
        let list1 = vec![("chunk-1".to_string(), 0.9), ("chunk-2".to_string(), 0.8)];
        let list2 = vec![("chunk-2".to_string(), 0.95), ("chunk-1".to_string(), 0.85)];

        let results = reciprocal_rank_fusion_scored(&[list1, list2], RRFConfig::default(), 10);

        assert_eq!(results.len(), 2);
        // Both chunks appear in both lists, RRF should fuse them
        assert!(results.iter().any(|r| r.chunk_id == "chunk-1"));
        assert!(results.iter().any(|r| r.chunk_id == "chunk-2"));
    }

    #[test]
    fn test_rrf_consensus_ranking() {
        // Test that RRF favors consensus
        let list1 = vec![
            "chunk-A".to_string(),
            "chunk-B".to_string(),
            "chunk-C".to_string(),
        ];
        let list2 = vec![
            "chunk-B".to_string(),
            "chunk-A".to_string(),
            "chunk-D".to_string(),
        ];
        let list3 = vec![
            "chunk-A".to_string(),
            "chunk-B".to_string(),
            "chunk-E".to_string(),
        ];

        let results = reciprocal_rank_fusion(&[list1, list2, list3], RRFConfig::default(), 5);

        // chunk-A and chunk-B appear in all lists, should rank highest
        assert_eq!(results[0].chunk_id, "chunk-A"); // Ranks 1, 2, 1 across lists
        assert_eq!(results[1].chunk_id, "chunk-B"); // Ranks 2, 1, 2 across
                                                    // lists
    }

    #[test]
    fn test_rrf_custom_k() {
        let list = vec!["chunk-1".to_string()];
        let config = RRFConfig { k: 10.0 };

        let results = reciprocal_rank_fusion(&[list], config, 10);

        // With k=10, score should be 1/(10+1) = 1/11 ≈ 0.0909
        let expected_score = 1.0 / 11.0;
        assert!((results[0].score - expected_score).abs() < 0.001);
    }
}
