src/evaluation/metrics.py

"""
Evaluation metrics for retrieval performance.

Computes F1@K, timing statistics, and per-modality breakdowns.
"""

import numpy as np
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass


@dataclass
class EvaluationResult:
    """Result of evaluation."""
    f1_at_5: float
    f1_at_10: float
    mean_time_ms: float
    median_time_ms: float
    p95_time_ms: float
    p99_time_ms: float
    modality_results: Dict[str, Dict[str, float]]


class EvaluationMetrics:
    """
    Evaluation metrics for retrieval systems.
    
    Computes F1@K and timing statistics.
    """
    
    def __init__(self):
        """Initialize evaluation metrics."""
        self._query_times: List[float] = []
    
    def compute_f1_at_k(
        self,
        predicted_indices: List[int],
        ground_truth_indices: List[int],
        k: int = 5
    ) -> float:
        """
        Compute F1@K between predicted and ground truth indices.
        
        Args:
            predicted_indices: Predicted indices (ranked)
            ground_truth_indices: Ground truth indices
            k: Top-K to consider
            
        Returns:
            F1 score (0-1)
        """
        # Take top-k predictions
        predicted_top_k = set(predicted_indices[:k])
        ground_truth_set = set(ground_truth_indices)
        
        # Compute precision and recall
        if len(predicted_top_k) == 0:
            precision = 0.0
        else:
            relevant_predicted = predicted_top_k & ground_truth_set
            precision = len(relevant_predicted) / len(predicted_top_k)
        
        if len(ground_truth_set) == 0:
            recall = 0.0
        else:
            relevant_predicted = predicted_top_k & ground_truth_set
            recall = len(relevant_predicted) / len(ground_truth_set)
        
        # Compute F1
        if precision + recall == 0:
            f1 = 0.0
        else:
            f1 = 2 * precision * recall / (precision + recall)
        
        return f1
    
    def compute_f1_at_k_batch(
        self,
        all_predicted: List[List[int]],
        all_ground_truth: List[List[int]],
        k: int = 5
    ) -> float:
        """
        Compute average F1@K over a batch.
        
        Args:
            all_predicted: List of predicted indices per query
            all_ground_truth: List of ground truth indices per query
            k: Top-K to consider
            
        Returns:
            Average F1 score
        """
        if len(all_predicted) == 0:
            return 0.0
        
        f1_scores = [
            self.compute_f1_at_k(pred, gt, k)
            for pred, gt in zip(all_predicted, all_ground_truth)
        ]
        
        return sum(f1_scores) / len(f1_scores)
    
    def add_query_time(self, time_ms: float) -> None:
        """Add a query time measurement."""
        self._query_times.append(time_ms)
    
    def get_timing_stats(self) -> Dict[str, float]:
        """
        Get timing statistics.
        
        Returns:
            Dict with mean, median, p95, p99
        """
        if not self._query_times:
            return {
                "mean_ms": 0.0,
                "median_ms": 0.0,
                "p95_ms": 0.0,
                "p99_ms": 0.0,
            }
        
        times = sorted(self._query_times)
        n = len(times)
        
        return {
            "mean_ms": sum(times) / n,
            "median_ms": times[n // 2],
            "p95_ms": times[int(n * 0.95)] if n >= 20 else times[-1],
            "p99_ms": times[int(n * 0.99)] if n >= 100 else times[-1],
        }
    
    def evaluate_retrieval(
        self,
        predicted_indices_list: List[List[int]],
        ground_truth_list: List[List[int]],
        query_times: Optional[List[float]] = None,
        modality_labels: Optional[List[str]] = None
    ) -> EvaluationResult:
        """
        Full evaluation of retrieval results.
        
        Args:
            predicted_indices_list: List of predicted indices per query
            ground_truth_list: List of ground truth indices per query
            query_times: Optional query times in ms
            modality_labels: Optional modality labels per query
            
        Returns:
            EvaluationResult with all metrics
        """
        # Compute F1@5 and F1@10
        f1_at_5 = self.compute_f1_at_k_batch(
            predicted_indices_list, ground_truth_list, k=5
        )
        f1_at_10 = self.compute_f1_at_k_batch(
            predicted_indices_list, ground_truth_list, k=10
        )
        
        # Timing stats
        if query_times:
            self._query_times.extend(query_times)
        
        timing = self.get_timing_stats()
        
        # Per-modality breakdown
        modality_results = {}
        if modality_labels:
            unique_modalities = set(modality_labels)
            
            for mod in unique_modalities:
                # Filter to this modality
                mod_indices = [
                    i for i, m in enumerate(modality_labels)
                    if m == mod
                ]
                
                if mod_indices:
                    mod_predicted = [predicted_indices_list[i] for i in mod_indices]
                    mod_gt = [ground_truth_list[i] for i in mod_indices]
                    
                    mod_f1_5 = self.compute_f1_at_k_batch(mod_predicted, mod_gt, k=5)
                    mod_f1_10 = self.compute_f1_at_k_batch(mod_predicted, mod_gt, k=10)
                    
                    modality_results[mod] = {
                        "f1_at_5": mod_f1_5,
                        "f1_at_10": mod_f1_10,
                        "n_queries": len(mod_indices),
                    }
        
        return EvaluationResult(
            f1_at_5=f1_at_5,
            f1_at_10=f1_at_10,
            mean_time_ms=timing["mean_ms"],
            median_time_ms=timing["median_ms"],
            p95_time_ms=timing["p95_ms"],
            p99_time_ms=timing["p99_ms"],
            modality_results=modality_results,
        )


# Self-check
if __name__ == "__main__":
    print("Testing EvaluationMetrics...")
    
    metrics = EvaluationMetrics()
    
    # Test F1 computation
    predicted = [0, 1, 2, 3, 4]
    ground_truth = [0, 2, 4, 6, 8]
    
    f1_5 = metrics.compute_f1_at_k(predicted, ground_truth, k=5)
    print(f"F1@5: {f1_5:.4f}")
    
    # Test batch F1
    all_predicted = [[0, 1, 2], [3, 4, 5]]
    all_gt = [[0, 1, 2], [3, 4, 5]]
    
    batch_f1 = metrics.compute_f1_at_k_batch(all_predicted, all_gt, k=3)
    print(f"Batch F1@3: {batch_f1:.4f}")
    
    # Test timing
    for t in [10.0, 20.0, 30.0, 40.0, 50.0]:
        metrics.add_query_time(t)
    
    timing = metrics.get_timing_stats()
    print(f"Timing stats: {timing}")
    
    # Test full evaluation
    result = metrics.evaluate_retrieval(
        all_predicted, all_gt,
        query_times=[15.0, 25.0],
        modality_labels=["optical", "sar"]
    )
    
    print(f"\nFull evaluation:")
    print(f"  F1@5: {result.f1_at_5:.4f}")
    print(f"  F1@10: {result.f1_at_10:.4f}")
    print(f"  Mean time: {result.mean_time_ms:.2f}ms")
    print(f"  Modality results: {result.modality_results}")
    
    print("\nEvaluationMetrics test passed!")