"""
Tests for multi-modal retrieval.

These tests verify:
- Same-modal retrieval (optical→optical, SAR→SAR, multispectral→multispectral)
- Cross-modal retrieval (optical↔SAR, optical↔multispectral)
- Modality filtering correctness
- Modality distribution tracking
"""

import pytest
import torch
import numpy as np

from src.retrieval.multimodal import MultiModalRetrieval, ModalityResult


@pytest.fixture
def multimodal_embeddings():
    """Create L2-normalized embeddings for all modalities."""
    n_per_modality = 50
    embed_dim = 768
    
    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),
    }
    
    # L2 normalize
    for mod in embeddings:
        embeddings[mod] = torch.nn.functional.normalize(embeddings[mod], dim=1)
    
    return embeddings


@pytest.fixture
def built_retrieval(multimodal_embeddings):
    """Create a MultiModalRetrieval with embeddings loaded."""
    retrieval = MultiModalRetrieval(embed_dim=768)
    retrieval.build_index(multimodal_embeddings)
    return retrieval


@pytest.fixture
def dummy_query():
    """Create a dummy query embedding."""
    query = torch.randn(768)
    return torch.nn.functional.normalize(query, dim=0)


class TestMultiModalRetrieval:
    """Tests for MultiModalRetrieval class."""
    
    def test_initialization(self):
        """Test initialization."""
        retrieval = MultiModalRetrieval(embed_dim=512)
        assert retrieval.embed_dim == 512
        assert retrieval.size == 0
    
    def test_build_index(self, multimodal_embeddings):
        """Test index builds correctly."""
        retrieval = MultiModalRetrieval(embed_dim=768)
        retrieval.build_index(multimodal_embeddings)
        
        assert retrieval.size == 150  # 50 * 3 modalities
        assert len(retrieval.modality_labels) == 150
    
    def test_modality_distribution(self, built_retrieval):
        """Test modality distribution tracking."""
        dist = built_retrieval.get_modality_distribution()
        
        assert dist["optical"] == 50
        assert dist["sar"] == 50
        assert dist["multispectral"] == 50
    
    def test_same_modal_optical(self, built_retrieval, dummy_query):
        """Test same-modal optical→optical retrieval."""
        result = built_retrieval.same_modal_query(
            dummy_query, modality="optical", k=5
        )
        
        assert isinstance(result, ModalityResult)
        assert len(result.indices) <= 5
        assert all(m == "optical" for m in result.modalities)
        assert result.query_modality == "optical"
    
    def test_same_modal_sar(self, built_retrieval, dummy_query):
        """Test same-modal SAR→SAR retrieval."""
        result = built_retrieval.same_modal_query(
            dummy_query, modality="sar", k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "sar" for m in result.modalities)
        assert result.query_modality == "sar"
    
    def test_same_modal_multispectral(self, built_retrieval, dummy_query):
        """Test same-modal multispectral→multispectral retrieval."""
        result = built_retrieval.same_modal_query(
            dummy_query, modality="multispectral", k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "multispectral" for m in result.modalities)
        assert result.query_modality == "multispectral"
    
    def test_cross_modal_optical_to_sar(self, built_retrieval, dummy_query):
        """Test cross-modal optical→SAR retrieval."""
        result = built_retrieval.cross_modal_query(
            dummy_query,
            source_modality="optical",
            target_modality="sar",
            k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "sar" for m in result.modalities)
        assert result.query_modality == "optical"
    
    def test_cross_modal_optical_to_multispectral(self, built_retrieval, dummy_query):
        """Test cross-modal optical→multispectral retrieval."""
        result = built_retrieval.cross_modal_query(
            dummy_query,
            source_modality="optical",
            target_modality="multispectral",
            k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "multispectral" for m in result.modalities)
    
    def test_cross_modal_sar_to_optical(self, built_retrieval, dummy_query):
        """Test cross-modal SAR→optical retrieval."""
        result = built_retrieval.cross_modal_query(
            dummy_query,
            source_modality="sar",
            target_modality="optical",
            k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "optical" for m in result.modalities)
    
    def test_cross_modal_sar_to_multispectral(self, built_retrieval, dummy_query):
        """Test cross-modal SAR→multispectral retrieval."""
        result = built_retrieval.cross_modal_query(
            dummy_query,
            source_modality="sar",
            target_modality="multispectral",
            k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "multispectral" for m in result.modalities)
    
    def test_cross_modal_multispectral_to_optical(self, built_retrieval, dummy_query):
        """Test cross-modal multispectral→optical retrieval."""
        result = built_retrieval.cross_modal_query(
            dummy_query,
            source_modality="multispectral",
            target_modality="optical",
            k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "optical" for m in result.modalities)
    
    def test_cross_modal_multispectral_to_sar(self, built_retrieval, dummy_query):
        """Test cross-modal multispectral→SAR retrieval."""
        result = built_retrieval.cross_modal_query(
            dummy_query,
            source_modality="multispectral",
            target_modality="sar",
            k=5
        )
        
        assert len(result.indices) <= 5
        assert all(m == "sar" for m in result.modalities)
    
    def test_mixed_query(self, built_retrieval, dummy_query):
        """Test mixed query returns all modalities."""
        result = built_retrieval.mixed_query(
            dummy_query, source_modality="optical", k=10
        )
        
        assert len(result.indices) <= 10
        assert len(result.modalities) == len(result.indices)
    
    def test_scores_valid(self, built_retrieval, dummy_query):
        """Test similarity scores are in valid range."""
        result = built_retrieval.same_modal_query(
            dummy_query, modality="optical", k=5
        )
        
        for score in result.scores:
            assert -1.0 <= score <= 1.0
    
    def test_indices_valid(self, built_retrieval, dummy_query):
        """Test returned indices are valid."""
        result = built_retrieval.same_modal_query(
            dummy_query, modality="optical", k=5
        )
        
        for idx in result.indices:
            assert 0 <= idx < built_retrieval.size


class TestModalityResult:
    """Tests for ModalityResult dataclass."""
    
    def test_modality_result_structure(self):
        """Test ModalityResult has correct fields."""
        result = ModalityResult(
            indices=[0, 1, 2],
            scores=[0.9, 0.8, 0.7],
            modalities=["optical", "optical", "optical"],
            query_modality="optical"
        )
        
        assert result.indices == [0, 1, 2]
        assert result.scores == [0.9, 0.8, 0.7]
        assert result.modalities == ["optical", "optical", "optical"]
        assert result.query_modality == "optical"


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