tests/test_competition_spec.py

"""
Competition Specification Tests for Cross-Modal Satellite Image Retrieval.

These tests verify the system meets ALL competition requirements:
- F1-score@5 for same-modal retrieval
- F1-score@10 for same-modal retrieval
- F1-score@5 for cross-modal retrieval
- F1-score@10 for cross-modal retrieval
- Average retrieval time per query
- Top-5 and Top-10 ranked results
- All modality combinations (optical, SAR, multispectral)
"""

import pytest
import torch
import time
import numpy as np
from typing import Dict, List, Tuple

from src.retrieval.multimodal import MultiModalRetrieval, ModalityResult
from src.retrieval.index import FAISSIndex
from src.evaluation.metrics import EvaluationMetrics, EvaluationResult
from src.evaluation.ground_truth import GroundTruth, create_ground_truth_from_matches


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def embed_dim():
    """Embedding dimension."""
    return 768


@pytest.fixture
def n_per_modality():
    """Number of samples per modality."""
    return 100


@pytest.fixture
def gallery_embeddings(embed_dim, n_per_modality):
    """
    Create L2-normalized gallery embeddings for all three modalities.
    Simulates a realistic gallery with paired samples.
    """
    torch.manual_seed(42)
    embeddings = {
        "optical": torch.randn(n_per_modality, embed_dim),
        "sar": torch.randn(n_per_modality, embed_dim),
        "multispectral": torch.randn(n_per_modality, embed_dim),
    }
    for mod in embeddings:
        embeddings[mod] = torch.nn.functional.normalize(embeddings[mod], dim=1)
    return embeddings


@pytest.fixture
def retrieval_index(gallery_embeddings, embed_dim):
    """
    Build a MultiModalRetrieval index from gallery embeddings.
    """
    retrieval = MultiModalRetrieval(embed_dim=embed_dim)
    retrieval.build_index(gallery_embeddings)
    return retrieval


@pytest.fixture
def modality_offsets():
    """
    FAISS index offsets for each modality in the combined index.
    Optical: 0..99, SAR: 100..199, Multispectral: 200..299
    """
    return {"optical": 0, "sar": 100, "multispectral": 200}


@pytest.fixture
def ground_truth_pairs(n_per_modality, modality_offsets):
    """
    Create ground-truth pairs for all same-modal and cross-modal combinations.

    FAISS index layout:
        optical        -> indices 0..99
        sar            -> indices 100..199
        multispectral  -> indices 200..299

    Sample i in modality A matches sample i in modality B, so the FAISS
    gallery ID is ``modality_offsets[modality] + i``.
    """
    matches = {}
    modalities = ["optical", "sar", "multispectral"]

    for src in modalities:
        for tgt in modalities:
            pairs = [
                (i, modality_offsets[tgt] + i)
                for i in range(n_per_modality)
            ]
            matches[(src, tgt)] = pairs

    return create_ground_truth_from_matches(matches)


@pytest.fixture
def evaluation_metrics():
    """Fresh EvaluationMetrics instance."""
    return EvaluationMetrics()


# ---------------------------------------------------------------------------
# Helper: run retrieval for all modality pairs and collect predictions
# ---------------------------------------------------------------------------

def run_all_queries(
    retrieval: MultiModalRetrieval,
    query_embeddings: Dict[str, torch.Tensor],
    k: int = 10,
) -> Tuple[Dict[str, List[List[int]]], Dict[str, List[float]]]:
    """
    Run queries for all same-modal and cross-modal combinations.
    
    Returns:
        predictions: dict mapping (src, tgt) -> list of predicted index lists
        query_times: dict mapping (src, tgt) -> list of query times in ms
    """
    modalities = ["optical", "sar", "multispectral"]
    predictions: Dict[str, List[List[int]]] = {}
    query_times: Dict[str, List[float]] = {}
    
    for src in modalities:
        for tgt in modalities:
            key = f"{src}_to_{tgt}"
            predictions[key] = []
            query_times[key] = []
            
            for query_emb in query_embeddings[src]:
                start = time.perf_counter()
                
                if src == tgt:
                    result = retrieval.same_modal_query(
                        query_emb.unsqueeze(0), modality=src, k=k
                    )
                else:
                    result = retrieval.cross_modal_query(
                        query_emb.unsqueeze(0),
                        source_modality=src,
                        target_modality=tgt,
                        k=k,
                    )
                
                elapsed_ms = (time.perf_counter() - start) * 1000
                predictions[key].append(result.indices)
                query_times[key].append(elapsed_ms)
    
    return predictions, query_times


# ===========================================================================
# TEST CLASS 1: Same-Modal Retrieval (F1@5 and F1@10)
# ===========================================================================

class TestSameModalRetrieval:
    """
    Competition Requirement:
        F1-score@5 for same-modal retrieval
        F1-score@10 for same-modal retrieval
    """
    
    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_returns_correct_modality(
        self, retrieval_index, modality, embed_dim
    ):
        """All returned results must belong to the queried modality."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality=modality, k=10)
        
        assert all(m == modality for m in result.modalities), (
            f"Same-modal query for {modality} returned wrong modality: {result.modalities}"
        )
    
    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_top5_count(self, retrieval_index, modality, embed_dim):
        """Same-modal query must return up to 5 results."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality=modality, k=5)
        
        assert len(result.indices) <= 5
        assert len(result.scores) == len(result.indices)
    
    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_top10_count(self, retrieval_index, modality, embed_dim):
        """Same-modal query must return up to 10 results."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality=modality, k=10)
        
        assert len(result.indices) <= 10
        assert len(result.scores) == len(result.indices)
    
    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_f1_at_5(
        self, retrieval_index, gallery_embeddings,
        evaluation_metrics, modality, embed_dim, n_per_modality, modality_offsets
    ):
        """
        F1@5 for same-modal retrieval must be computable and >= 0.
        """
        n_queries = 20
        offset = modality_offsets[modality]
        query_embs = gallery_embeddings[modality][:n_queries]

        all_predicted = []
        all_gt = []

        for i in range(n_queries):
            result = retrieval_index.same_modal_query(
                query_embs[i].unsqueeze(0), modality=modality, k=5
            )
            all_predicted.append(result.indices)
            # Ground truth: sample i matches FAISS index (offset + i)
            all_gt.append([offset + i])

        f1_5 = evaluation_metrics.compute_f1_at_k_batch(all_predicted, all_gt, k=5)

        assert 0.0 <= f1_5 <= 1.0, f"F1@5 out of range: {f1_5}"
        assert f1_5 > 0.0, f"F1@5 should be > 0 for same-modal {modality}"

    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_f1_at_10(
        self, retrieval_index, gallery_embeddings,
        evaluation_metrics, modality, embed_dim, n_per_modality, modality_offsets
    ):
        """
        F1@10 for same-modal retrieval must be computable and >= 0.
        """
        n_queries = 20
        offset = modality_offsets[modality]
        query_embs = gallery_embeddings[modality][:n_queries]

        all_predicted = []
        all_gt = []

        for i in range(n_queries):
            result = retrieval_index.same_modal_query(
                query_embs[i].unsqueeze(0), modality=modality, k=10
            )
            all_predicted.append(result.indices)
            all_gt.append([offset + i])

        f1_10 = evaluation_metrics.compute_f1_at_k_batch(all_predicted, all_gt, k=10)

        assert 0.0 <= f1_10 <= 1.0, f"F1@10 out of range: {f1_10}"
        assert f1_10 > 0.0, f"F1@10 should be > 0 for same-modal {modality}"


# ===========================================================================
# TEST CLASS 2: Cross-Modal Retrieval (F1@5 and F1@10)
# ===========================================================================

class TestCrossModalRetrieval:
    """
    Competition Requirement:
        F1-score@5 for cross-modal retrieval
        F1-score@10 for cross-modal retrieval
        (optical->SAR, SAR->optical, optical->multispectral, etc.)
    """
    
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("optical", "multispectral"),
        ("sar", "optical"),
        ("sar", "multispectral"),
        ("multispectral", "optical"),
        ("multispectral", "sar"),
    ])
    def test_cross_modal_returns_target_modality(
        self, retrieval_index, source, target, embed_dim
    ):
        """Cross-modal query must only return results from the target modality."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.cross_modal_query(
            query, source_modality=source, target_modality=target, k=10
        )
        
        assert all(m == target for m in result.modalities), (
            f"Cross-modal {source}->{target} returned wrong modality: {result.modalities}"
        )
        assert result.query_modality == source
    
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("optical", "multispectral"),
        ("sar", "optical"),
        ("sar", "multispectral"),
        ("multispectral", "optical"),
        ("multispectral", "sar"),
    ])
    def test_cross_modal_top5_count(
        self, retrieval_index, source, target, embed_dim
    ):
        """Cross-modal query must return up to 5 results."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.cross_modal_query(
            query, source_modality=source, target_modality=target, k=5
        )
        
        assert len(result.indices) <= 5
    
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("optical", "multispectral"),
        ("sar", "optical"),
        ("sar", "multispectral"),
        ("multispectral", "optical"),
        ("multispectral", "sar"),
    ])
    def test_cross_modal_top10_count(
        self, retrieval_index, source, target, embed_dim
    ):
        """Cross-modal query must return up to 10 results."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.cross_modal_query(
            query, source_modality=source, target_modality=target, k=10
        )
        
        assert len(result.indices) <= 10
    
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("optical", "multispectral"),
        ("sar", "optical"),
        ("sar", "multispectral"),
        ("multispectral", "optical"),
        ("multispectral", "sar"),
    ])
    def test_cross_modal_f1_at_5(
        self, retrieval_index, gallery_embeddings, ground_truth_pairs,
        evaluation_metrics, source, target, embed_dim, n_per_modality
    ):
        """
        F1@5 for cross-modal retrieval must be computable and >= 0.
        """
        n_queries = 20
        query_embs = gallery_embeddings[source][:n_queries]
        
        all_predicted = []
        all_gt = []
        
        for i in range(n_queries):
            result = retrieval_index.cross_modal_query(
                query_embs[i].unsqueeze(0),
                source_modality=source,
                target_modality=target,
                k=5,
            )
            all_predicted.append(result.indices)
            gt_ids = ground_truth_pairs.get_gallery_ids(i)
            all_gt.append(gt_ids)
        
        f1_5 = evaluation_metrics.compute_f1_at_k_batch(all_predicted, all_gt, k=5)
        
        assert 0.0 <= f1_5 <= 1.0, f"F1@5 out of range for {source}->{target}: {f1_5}"
    
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("optical", "multispectral"),
        ("sar", "optical"),
        ("sar", "multispectral"),
        ("multispectral", "optical"),
        ("multispectral", "sar"),
    ])
    def test_cross_modal_f1_at_10(
        self, retrieval_index, gallery_embeddings, ground_truth_pairs,
        evaluation_metrics, source, target, embed_dim, n_per_modality
    ):
        """
        F1@10 for cross-modal retrieval must be computable and >= 0.
        """
        n_queries = 20
        query_embs = gallery_embeddings[source][:n_queries]
        
        all_predicted = []
        all_gt = []
        
        for i in range(n_queries):
            result = retrieval_index.cross_modal_query(
                query_embs[i].unsqueeze(0),
                source_modality=source,
                target_modality=target,
                k=10,
            )
            all_predicted.append(result.indices)
            gt_ids = ground_truth_pairs.get_gallery_ids(i)
            all_gt.append(gt_ids)
        
        f1_10 = evaluation_metrics.compute_f1_at_k_batch(all_predicted, all_gt, k=10)
        
        assert 0.0 <= f1_10 <= 1.0, f"F1@10 out of range for {source}->{target}: {f1_10}"


# ===========================================================================
# TEST CLASS 3: Retrieval Time Measurement
# ===========================================================================

class TestRetrievalTime:
    """
    Competition Requirement:
        Average retrieval time per query (including feature matching
        and similarity search time).
    """
    
    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_retrieval_time_recorded(
        self, retrieval_index, modality, embed_dim
    ):
        """Each same-modal query must record a positive retrieval time."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        start = time.perf_counter()
        retrieval_index.same_modal_query(query, modality=modality, k=5)
        elapsed_ms = (time.perf_counter() - start) * 1000
        
        assert elapsed_ms > 0, "Retrieval time must be positive"
    
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("sar", "optical"),
        ("optical", "multispectral"),
    ])
    def test_cross_modal_retrieval_time_recorded(
        self, retrieval_index, source, target, embed_dim
    ):
        """Each cross-modal query must record a positive retrieval time."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        start = time.perf_counter()
        retrieval_index.cross_modal_query(
            query, source_modality=source, target_modality=target, k=5
        )
        elapsed_ms = (time.perf_counter() - start) * 1000
        
        assert elapsed_ms > 0, "Cross-modal retrieval time must be positive"
    
    def test_average_retrieval_time_benchmark(
        self, retrieval_index, embed_dim
    ):
        """
        Average retrieval time per query must be measured.
        Competition prefers lower retrieval time.
        """
        n_queries = 50
        times = []
        
        for _ in range(n_queries):
            query = torch.randn(1, embed_dim)
            query = torch.nn.functional.normalize(query, dim=1)
            
            start = time.perf_counter()
            retrieval_index.same_modal_query(query, modality="optical", k=5)
            elapsed_ms = (time.perf_counter() - start) * 1000
            times.append(elapsed_ms)
        
        avg_time = sum(times) / len(times)
        
        assert avg_time > 0, "Average retrieval time must be positive"
        # Competition-grade: average should be under 1 second
        assert avg_time < 1000, f"Average retrieval time too high: {avg_time:.2f}ms"
    
    def test_retrieval_time_stats(
        self, retrieval_index, evaluation_metrics, embed_dim
    ):
        """Timing statistics (mean, median, p95, p99) must be computable."""
        n_queries = 30
        times = []
        
        for _ in range(n_queries):
            query = torch.randn(1, embed_dim)
            query = torch.nn.functional.normalize(query, dim=1)
            
            start = time.perf_counter()
            retrieval_index.same_modal_query(query, modality="optical", k=10)
            elapsed_ms = (time.perf_counter() - start) * 1000
            times.append(elapsed_ms)
            evaluation_metrics.add_query_time(elapsed_ms)
        
        stats = evaluation_metrics.get_timing_stats()
        
        assert stats["mean_ms"] > 0
        assert stats["median_ms"] > 0
        assert stats["p95_ms"] > 0
        # p99 may equal p95 with small samples, just check it exists
        assert "p99_ms" in stats


# ===========================================================================
# TEST CLASS 4: Top-5 and Top-10 Ranking
# ===========================================================================

class TestRanking:
    """
    Competition Requirement:
        Ranking of the top-5 and top-10 most relevant images.
        Results must be sorted by similarity score (descending).
    """
    
    @pytest.mark.parametrize("k", [5, 10])
    @pytest.mark.parametrize("modality", ["optical", "sar", "multispectral"])
    def test_same_modal_topk_sorted_descending(
        self, retrieval_index, modality, embed_dim, k
    ):
        """Top-K results must be sorted by score in descending order."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality=modality, k=k)
        
        scores = result.scores
        for i in range(len(scores) - 1):
            assert scores[i] >= scores[i + 1], (
                f"Scores not sorted descending at position {i}: "
                f"{scores[i]} < {scores[i+1]}"
            )
    
    @pytest.mark.parametrize("k", [5, 10])
    @pytest.mark.parametrize("source,target", [
        ("optical", "sar"),
        ("sar", "optical"),
        ("multispectral", "sar"),
    ])
    def test_cross_modal_topk_sorted_descending(
        self, retrieval_index, source, target, embed_dim, k
    ):
        """Cross-modal Top-K results must be sorted by score descending."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.cross_modal_query(
            query, source_modality=source, target_modality=target, k=k
        )
        
        scores = result.scores
        for i in range(len(scores) - 1):
            assert scores[i] >= scores[i + 1], (
                f"Cross-modal scores not sorted for {source}->{target} at position {i}"
            )
    
    @pytest.mark.parametrize("k", [5, 10])
    def test_same_modal_topk_unique_indices(
        self, retrieval_index, embed_dim, k
    ):
        """Top-K results must have unique indices (no duplicates)."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality="optical", k=k)
        
        assert len(result.indices) == len(set(result.indices)), (
            f"Duplicate indices in top-{k} results: {result.indices}"
        )
    
    @pytest.mark.parametrize("k", [5, 10])
    def test_same_modal_topk_score_range(
        self, retrieval_index, embed_dim, k
    ):
        """Scores must be valid cosine similarity values (-1 to 1)."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality="optical", k=k)
        
        for score in result.scores:
            assert -1.0 <= score <= 1.0, f"Score out of range: {score}"


# ===========================================================================
# TEST CLASS 5: End-to-End Integration
# ===========================================================================

class TestEndToEndIntegration:
    """
    Full pipeline test: build index -> query -> evaluate -> report metrics.
    Simulates the competition evaluation flow.
    """
    
    def test_full_evaluation_pipeline(
        self, retrieval_index, gallery_embeddings, ground_truth_pairs,
        evaluation_metrics, embed_dim, n_per_modality
    ):
        """
        Complete evaluation pipeline producing F1@5, F1@10, and timing stats
        for all modality combinations.
        """
        modalities = ["optical", "sar", "multispectral"]
        n_queries = 10
        
        all_predictions = []
        all_ground_truth = []
        all_times = []
        all_modality_labels = []
        
        for src in modalities:
            for tgt in modalities:
                query_embs = gallery_embeddings[src][:n_queries]
                
                for i in range(n_queries):
                    start = time.perf_counter()
                    
                    if src == tgt:
                        result = retrieval_index.same_modal_query(
                            query_embs[i].unsqueeze(0), modality=src, k=10
                        )
                    else:
                        result = retrieval_index.cross_modal_query(
                            query_embs[i].unsqueeze(0),
                            source_modality=src,
                            target_modality=tgt,
                            k=10,
                        )
                    
                    elapsed_ms = (time.perf_counter() - start) * 1000
                    
                    all_predictions.append(result.indices)
                    all_ground_truth.append(ground_truth_pairs.get_gallery_ids(i))
                    all_times.append(elapsed_ms)
                    all_modality_labels.append(f"{src}_to_{tgt}")
        
        # Run evaluation
        eval_result = evaluation_metrics.evaluate_retrieval(
            all_predictions,
            all_ground_truth,
            query_times=all_times,
            modality_labels=all_modality_labels,
        )
        
        # Verify all required metrics exist
        assert hasattr(eval_result, "f1_at_5")
        assert hasattr(eval_result, "f1_at_10")
        assert hasattr(eval_result, "mean_time_ms")
        assert hasattr(eval_result, "median_time_ms")
        assert hasattr(eval_result, "p95_time_ms")
        assert hasattr(eval_result, "p99_time_ms")
        assert hasattr(eval_result, "modality_results")
        
        # Verify metric ranges
        assert 0.0 <= eval_result.f1_at_5 <= 1.0
        assert 0.0 <= eval_result.f1_at_10 <= 1.0
        assert eval_result.mean_time_ms > 0
        assert eval_result.median_time_ms > 0
        
        # Verify per-modality breakdown exists
        assert len(eval_result.modality_results) > 0
    
    def test_all_modality_pairs_covered(
        self, retrieval_index, gallery_embeddings, embed_dim
    ):
        """
        All 9 modality combinations (3x3) must be queryable.
        """
        modalities = ["optical", "sar", "multispectral"]
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        for src in modalities:
            for tgt in modalities:
                if src == tgt:
                    result = retrieval_index.same_modal_query(
                        query, modality=src, k=5
                    )
                else:
                    result = retrieval_index.cross_modal_query(
                        query, source_modality=src, target_modality=tgt, k=5
                    )
                
                assert len(result.indices) > 0, (
                    f"No results for {src}->{tgt}"
                )
    
    def test_competition_report_format(
        self, retrieval_index, gallery_embeddings, ground_truth_pairs,
        evaluation_metrics, embed_dim, n_per_modality
    ):
        """
        Generate a competition-style report with all required metrics.
        """
        modalities = ["optical", "sar", "multispectral"]
        n_queries = 10
        
        report = {}
        
        for src in modalities:
            for tgt in modalities:
                key = f"{src}_to_{tgt}"
                query_embs = gallery_embeddings[src][:n_queries]
                
                predictions_5 = []
                predictions_10 = []
                ground_truth_list = []
                times = []
                
                for i in range(n_queries):
                    start = time.perf_counter()
                    
                    if src == tgt:
                        result_5 = retrieval_index.same_modal_query(
                            query_embs[i].unsqueeze(0), modality=src, k=5
                        )
                        result_10 = retrieval_index.same_modal_query(
                            query_embs[i].unsqueeze(0), modality=src, k=10
                        )
                    else:
                        result_5 = retrieval_index.cross_modal_query(
                            query_embs[i].unsqueeze(0),
                            source_modality=src, target_modality=tgt, k=5
                        )
                        result_10 = retrieval_index.cross_modal_query(
                            query_embs[i].unsqueeze(0),
                            source_modality=src, target_modality=tgt, k=10
                        )
                    
                    elapsed_ms = (time.perf_counter() - start) * 1000
                    
                    predictions_5.append(result_5.indices)
                    predictions_10.append(result_10.indices)
                    ground_truth_list.append(ground_truth_pairs.get_gallery_ids(i))
                    times.append(elapsed_ms)
                
                f1_5 = evaluation_metrics.compute_f1_at_k_batch(
                    predictions_5, ground_truth_list, k=5
                )
                f1_10 = evaluation_metrics.compute_f1_at_k_batch(
                    predictions_10, ground_truth_list, k=10
                )
                avg_time = sum(times) / len(times)
                
                report[key] = {
                    "f1_at_5": f1_5,
                    "f1_at_10": f1_10,
                    "avg_time_ms": avg_time,
                    "n_queries": n_queries,
                }
        
        # Verify report structure
        assert len(report) == 9  # 3x3 modality combinations
        
        for key, metrics in report.items():
            assert "f1_at_5" in metrics
            assert "f1_at_10" in metrics
            assert "avg_time_ms" in metrics
            assert "n_queries" in metrics
            assert 0.0 <= metrics["f1_at_5"] <= 1.0
            assert 0.0 <= metrics["f1_at_10"] <= 1.0
            assert metrics["avg_time_ms"] > 0


# ===========================================================================
# TEST CLASS 6: Edge Cases and Robustness
# ===========================================================================

class TestEdgeCases:
    """Robustness tests for the retrieval system."""
    
    def test_single_query_embedding(self, retrieval_index, embed_dim):
        """System must handle a single query vector."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality="optical", k=5)
        
        assert len(result.indices) > 0
    
    def test_k_larger_than_gallery(self, retrieval_index, embed_dim):
        """k larger than gallery size should not crash."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        # k=200 but gallery has only 100 per modality
        result = retrieval_index.same_modal_query(query, modality="optical", k=200)
        
        assert len(result.indices) <= 100  # clamped to gallery size
    
    def test_k_equals_one(self, retrieval_index, embed_dim):
        """k=1 must return exactly one result."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality="optical", k=1)
        
        assert len(result.indices) == 1
        assert len(result.scores) == 1
    
    def test_repeated_queries_consistent(self, retrieval_index, embed_dim):
        """Same query repeated must return identical results."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        r1 = retrieval_index.same_modal_query(query, modality="optical", k=5)
        r2 = retrieval_index.same_modal_query(query, modality="optical", k=5)
        
        assert r1.indices == r2.indices
        assert r1.scores == r2.scores
    
    def test_score_non_negative_for_normalized(
        self, retrieval_index, embed_dim
    ):
        """Cosine similarity for L2-normalized vectors should be >= 0
        when vectors are random (most pairs have positive similarity)."""
        query = torch.randn(1, embed_dim)
        query = torch.nn.functional.normalize(query, dim=1)
        
        result = retrieval_index.same_modal_query(query, modality="optical", k=5)
        
        # With random normalized vectors, scores are typically positive
        # Just verify they are valid
        for score in result.scores:
            assert -1.0 <= score <= 1.0


# ---------------------------------------------------------------------------
# Self-check
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])