ranking_metrics.py

"""
Ranking Quality Metrics

Implements standard information retrieval metrics for evaluating policy recommendation rankings:
- NDCG@k: Normalized Discounted Cumulative Gain at position k
- Precision@k: Fraction of top-k results that are relevant
- Mean Average Precision (MAP): Average of precision values across recall levels

These are generic, reusable metrics for any ranking task.
"""

import numpy as np
from typing import List, Dict, Tuple


def dcg_at_k(relevances: List[float], k: int) -> float:
    """
    Calculate Discounted Cumulative Gain at position k.
    
    DCG = sum(rel_i / log2(i+1)) for i in 0..k-1
    
    Args:
        relevances: List of relevance scores (0-1) for items in ranked order
        k: Position cutoff
    
    Returns:
        DCG value (higher is better)
    """
    if not relevances or k <= 0:
        return 0.0
    
    dcg = 0.0
    for i in range(min(k, len(relevances))):
        dcg += relevances[i] / np.log2(i + 2)  # i+2 because log2(1)=0, we want log2(2)=1 for first item
    
    return dcg


def idcg_at_k(relevances: List[float], k: int) -> float:
    """
    Calculate Ideal Discounted Cumulative Gain at position k.
    
    IDCG is the DCG of the perfect ranking (sorted in descending order).
    
    Args:
        relevances: List of all possible relevance scores
        k: Position cutoff
    
    Returns:
        IDCG value
    """
    if not relevances or k <= 0:
        return 0.0
    
    # Sort in descending order (best items first)
    sorted_relevances = sorted(relevances, reverse=True)
    
    return dcg_at_k(sorted_relevances, k)


def ndcg_at_k(relevances: List[float], k: int) -> float:
    """
    Calculate Normalized Discounted Cumulative Gain at position k.
    
    NDCG@k = DCG@k / IDCG@k
    
    Ranges from 0 to 1, where 1 is perfect ranking.
    
    Args:
        relevances: List of relevance scores (0-1) for items in ranked order
        k: Position cutoff (e.g., 5 for NDCG@5)
    
    Returns:
        NDCG@k value (0-1)
    """
    if not relevances:
        return 0.0
    
    dcg = dcg_at_k(relevances, k)
    idcg = idcg_at_k(relevances, k)
    
    if idcg == 0.0:
        return 0.0
    
    return dcg / idcg


def precision_at_k(relevances: List[float], k: int, threshold: float = 0.5) -> float:
    """
    Calculate Precision at position k.
    
    Precision@k = (# relevant items in top-k) / k
    
    An item is considered relevant if relevance >= threshold (default 0.5).
    
    Args:
        relevances: List of relevance scores (0-1) for items in ranked order
        k: Position cutoff
        threshold: Relevance threshold for considering an item relevant
    
    Returns:
        Precision@k (0-1)
    """
    if not relevances or k <= 0:
        return 0.0
    
    top_k = relevances[:k]
    relevant_count = sum(1 for rel in top_k if rel >= threshold)
    
    return relevant_count / k


def mean_average_precision(relevances: List[List[float]], k: int = 10) -> float:
    """
    Calculate Mean Average Precision (MAP).
    
    MAP averages precision values across all recall levels for multiple queries/rankings.
    
    Args:
        relevances: List of relevance score lists, one per query/ranking
        k: Maximum position cutoff to consider
    
    Returns:
        MAP value (0-1)
    """
    if not relevances:
        return 0.0
    
    average_precisions = []
    
    for rel_list in relevances:
        if not rel_list:
            average_precisions.append(0.0)
            continue
        
        # Calculate precision at each position
        precisions = []
        relevant_count = 0
        
        for i in range(min(k, len(rel_list))):
            # Typically threshold is > 0 for relevance
            if rel_list[i] > 0.0:
                relevant_count += 1
                precision_at_i = relevant_count / (i + 1)
                precisions.append(precision_at_i)
        
        # Average precision for this query
        if precisions:
            ap = sum(precisions) / len(precisions)
        else:
            ap = 0.0
        
        average_precisions.append(ap)
    
    # Mean across all queries
    return sum(average_precisions) / len(average_precisions) if average_precisions else 0.0


def reciprocal_rank(relevances: List[float], threshold: float = 0.5) -> float:
    """
    Calculate Reciprocal Rank (MRR component).
    
    Reciprocal Rank = 1 / rank of first relevant item (or 0 if none)
    
    Args:
        relevances: List of relevance scores in ranked order
        threshold: Relevance threshold
    
    Returns:
        Reciprocal rank (0-1)
    """
    for i, rel in enumerate(relevances):
        if rel >= threshold:
            return 1.0 / (i + 1)
    
    return 0.0


def mean_reciprocal_rank(relevances_list: List[List[float]], threshold: float = 0.5) -> float:
    """
    Calculate Mean Reciprocal Rank (MRR).
    
    MRR = mean of reciprocal ranks across multiple queries/rankings.
    
    Args:
        relevances_list: List of relevance score lists
        threshold: Relevance threshold
    
    Returns:
        MRR value (0-1)
    """
    if not relevances_list:
        return 0.0
    
    rrs = [reciprocal_rank(rel, threshold) for rel in relevances_list]
    
    return sum(rrs) / len(rrs)


def evaluate_ranking(
    relevances: List[float],
    k: int = 5,
    threshold: float = 0.5
) -> Dict[str, float]:
    """
    Comprehensive ranking evaluation.
    
    Computes multiple metrics for a single ranked list.
    
    Args:
        relevances: List of relevance scores in ranked order
        k: Position cutoff
        threshold: Relevance threshold for Precision@k
    
    Returns:
        Dictionary with metrics: ndcg@k, precision@k, map, mrr
    """
    return {
        f"ndcg@{k}": ndcg_at_k(relevances, k),
        f"precision@{k}": precision_at_k(relevances, k, threshold),
        "mrr": reciprocal_rank(relevances, threshold),
    }


def evaluate_ranking_batch(
    relevances_list: List[List[float]],
    k: int = 5,
    threshold: float = 0.5
) -> Dict[str, float]:
    """
    Comprehensive ranking evaluation across multiple rankings.
    
    Args:
        relevances_list: List of relevance score lists
        k: Position cutoff
        threshold: Relevance threshold
    
    Returns:
        Dictionary with aggregated metrics: ndcg@k, precision@k, map, mrr
    """
    ndcg_values = [ndcg_at_k(rel, k) for rel in relevances_list]
    precision_values = [precision_at_k(rel, k, threshold) for rel in relevances_list]
    
    return {
        f"ndcg@{k}": np.mean(ndcg_values) if ndcg_values else 0.0,
        f"precision@{k}": np.mean(precision_values) if precision_values else 0.0,
        "map": mean_average_precision(relevances_list, k),
        "mrr": mean_reciprocal_rank(relevances_list, threshold),
    }