"""

Mean Reciprocal Rank (MRR) metric

"""

import numpy as np

from typing import Dict, List, Tuple

def compute_mrr_at_k(

qrels: Dict[str, Dict[str, int]],

results: Dict[str, List[Tuple[str, int, float]]],

k: int = 10,

) -> float:

"""

Compute Mean Reciprocal Rank at k.

MRR@k measures the average of the reciprocal ranks of the first relevant

document found within the top-k results for each query.

Args:

qrels: Dictionary mapping query_id to {passage_id: relevance}

results: Dictionary mapping query_id to [(passage_id, rank, score), ...]

k: Cutoff rank (default: 10)

Returns:

float: MRR@k score (0.0 to 1.0)

Example:

If the first relevant document appears at rank 3, RR = 1/3 = 0.333

MRR is the average of RR across all queries.

"""

reciprocal_ranks = []

for query_id, query_qrels in qrels.items():

if query_id not in results:

reciprocal_ranks.append(0.0)

continue

# Get top-k results

top_k_results = results[query_id][:k]

# Find first relevant document

for passage_id, rank, _ in top_k_results:

if passage_id in query_qrels and query_qrels[passage_id] > 0:

reciprocal_ranks.append(1.0 / rank)

break

else:

# No relevant document found in top-k

reciprocal_ranks.append(0.0)

return np.mean(reciprocal_ranks) if reciprocal_ranks else 0.0