//! Retrieval metrics: recall@k, precision@k, MRR, nDCG with graded
//! relevance. Includes the [`RelevanceJudgment`] type that serves both as
//! the metric's input contract and as the standard golden-data row shape
//! for any IR evaluation harness.

use std::collections::{HashMap, HashSet};

use serde::{Deserialize, Serialize};

/// A relevance judgment: a document ID and its graded relevance.
///
/// The grade convention is standard for graded-relevance IR evaluation:
/// `0` = not relevant, `1` = marginally, `2` = relevant, `3` = highly
/// relevant. The kernel treats any positive grade as a "relevant" hit for
/// recall / precision / MRR and uses the raw grade for nDCG's ideal-DCG
/// computation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelevanceJudgment {
    pub doc_id: String,
    pub grade: i32,
}

/// Per-query metric values.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct QueryMetrics {
    pub recall: f64,
    pub precision: f64,
    pub mrr: f64,
    pub ndcg: f64,
}

/// Aggregate metrics across all queries (mean).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AggregateMetrics {
    pub recall_at_k: f64,
    pub precision_at_k: f64,
    pub mrr: f64,
    pub ndcg: f64,
}

pub struct RetrievalMetrics;

impl RetrievalMetrics {
    /// Compute metrics for a single query given retrieved IDs and relevance judgments.
    pub fn compute_query(
        retrieved_ids: &[String],
        judgments: &[RelevanceJudgment],
        k: usize,
    ) -> QueryMetrics {
        let top_k = &retrieved_ids[..k.min(retrieved_ids.len())];

        let relevant_set: HashSet<&str> = judgments
            .iter()
            .filter(|j| j.grade > 0)
            .map(|j| j.doc_id.as_str())
            .collect();

        let hits: Vec<bool> = top_k
            .iter()
            .map(|id| relevant_set.contains(id.as_str()))
            .collect();

        // Recall@k
        let recall = if relevant_set.is_empty() {
            0.0
        } else {
            hits.iter().filter(|&&h| h).count() as f64 / relevant_set.len() as f64
        };

        // Precision@k
        let precision = if k == 0 {
            0.0
        } else {
            hits.iter().filter(|&&h| h).count() as f64 / k as f64
        };

        // MRR (Mean Reciprocal Rank)
        let mrr = hits
            .iter()
            .position(|&h| h)
            .map(|pos| 1.0 / (pos + 1) as f64)
            .unwrap_or(0.0);

        // nDCG with graded relevance
        let grade_map: HashMap<&str, i32> = judgments
            .iter()
            .map(|j| (j.doc_id.as_str(), j.grade))
            .collect();

        let dcg: f64 = top_k
            .iter()
            .enumerate()
            .map(|(i, id)| {
                let gain = *grade_map.get(id.as_str()).unwrap_or(&0) as f64;
                gain / (i as f64 + 2.0).log2()
            })
            .sum();

        let mut ideal_gains: Vec<f64> = judgments.iter().map(|j| j.grade as f64).collect();
        ideal_gains.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
        let ideal_dcg: f64 = ideal_gains
            .iter()
            .take(k)
            .enumerate()
            .map(|(i, &gain)| gain / (i as f64 + 2.0).log2())
            .sum();

        let ndcg = if ideal_dcg > 0.0 {
            dcg / ideal_dcg
        } else {
            0.0
        };

        QueryMetrics {
            recall,
            precision,
            mrr,
            ndcg,
        }
    }

    /// Recall@k for each k in `ks`, for a single query.
    ///
    /// Returns one `(k, recall@k)` pair per requested cutoff, in the order
    /// `ks` was given. This is the multi-cutoff companion to
    /// [`Self::compute_query`] (which records recall at a single k): it reuses
    /// the exact same recall definition — `|relevant ∩ top-k| / |relevant|`,
    /// `0.0` when no judgments are positive — applied at several cutoffs so a
    /// per-query record can carry Recall@{1,3,5,10} without re-deriving the
    /// formula. Each cutoff is clamped to the number of retrieved ids, matching
    /// `compute_query`'s `k.min(len)` slicing.
    pub fn recall_at_ks(
        retrieved_ids: &[String],
        judgments: &[RelevanceJudgment],
        ks: &[usize],
    ) -> Vec<(usize, f64)> {
        let relevant_set: HashSet<&str> = judgments
            .iter()
            .filter(|j| j.grade > 0)
            .map(|j| j.doc_id.as_str())
            .collect();

        ks.iter()
            .map(|&k| {
                let recall = if relevant_set.is_empty() {
                    0.0
                } else {
                    let top_k = &retrieved_ids[..k.min(retrieved_ids.len())];
                    let hits = top_k
                        .iter()
                        .filter(|id| relevant_set.contains(id.as_str()))
                        .count();
                    hits as f64 / relevant_set.len() as f64
                };
                (k, recall)
            })
            .collect()
    }

    /// Average metrics across multiple queries.
    pub fn aggregate(queries: &[QueryMetrics]) -> AggregateMetrics {
        let n = queries.len() as f64;
        if n == 0.0 {
            return AggregateMetrics::default();
        }
        AggregateMetrics {
            recall_at_k: queries.iter().map(|q| q.recall).sum::<f64>() / n,
            precision_at_k: queries.iter().map(|q| q.precision).sum::<f64>() / n,
            mrr: queries.iter().map(|q| q.mrr).sum::<f64>() / n,
            ndcg: queries.iter().map(|q| q.ndcg).sum::<f64>() / n,
        }
    }
}
