src/retrieval/multimodal.py

"""
Multi-modal retrieval for satellite imagery.

Handles same-modal and cross-modal retrieval with modality filtering.
"""

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

from .index import FAISSIndex


@dataclass
class ModalityResult:
    """Result with modality information."""
    indices: List[int]
    scores: List[float]
    modalities: List[str]
    query_modality: str


class MultiModalRetrieval:
    """
    Multi-modal retrieval with modality-aware search.
    
    Supports same-modal and cross-modal queries with filtering.
    """
    
    # Modality to index mapping
    MODALITY_MAP = {
        "optical": 0,
        "sar": 1,
        "multispectral": 2,
    }
    
    def __init__(self, embed_dim: int = 768):
        """
        Initialize multi-modal retrieval.
        
        Args:
            embed_dim: Embedding dimension
        """
        self.embed_dim = embed_dim
        self.index = FAISSIndex(embed_dim)
        
        # Track modality for each embedding
        self.modality_labels: List[str] = []
        self.sample_ids: List[int] = []
    
    @property
    def size(self) -> int:
        """Total number of embeddings."""
        return self.index.size
    
    def build_index(
        self,
        embeddings_by_modality: Dict[str, torch.Tensor],
        sample_ids_by_modality: Optional[Dict[str, List[int]]] = None
    ) -> None:
        """
        Build index with modality labels.
        
        Args:
            embeddings_by_modality: Dict mapping modality to embeddings tensor
            sample_ids_by_modality: Optional sample IDs per modality
        """
        all_embeddings = []
        all_modalities = []
        all_sample_ids = []
        
        for modality, embeddings in embeddings_by_modality.items():
            # Convert to numpy if needed
            if isinstance(embeddings, torch.Tensor):
                embeddings = embeddings.numpy().astype(np.float32)
            
            all_embeddings.append(embeddings)
            all_modalities.extend([modality] * len(embeddings))
            
            # Sample IDs
            if sample_ids_by_modality and modality in sample_ids_by_modality:
                all_sample_ids.extend(sample_ids_by_modality[modality])
            else:
                all_sample_ids.extend(range(len(embeddings)))
        
        # Concatenate all embeddings
        combined_embeddings = np.concatenate(all_embeddings, axis=0)
        
        # Build index
        self.index.build(combined_embeddings)
        self.modality_labels = all_modalities
        self.sample_ids = all_sample_ids
    
    def _filter_by_modality(
        self,
        indices: np.ndarray,
        scores: np.ndarray,
        target_modality: Optional[str] = None
    ) -> Tuple[List[int], List[float], List[str]]:
        """
        Filter results by modality.
        
        Args:
            indices: Raw indices from FAISS
            scores: Raw scores from FAISS
            target_modality: If specified, only return results from this modality
            
        Returns:
            (filtered_indices, filtered_scores, modalities)
        """
        filtered_indices = []
        filtered_scores = []
        filtered_modalities = []
        
        for idx, score in zip(indices[0], scores[0]):
            if idx < 0:  # FAISS returns -1 for empty slots
                continue
            
            modality = self.modality_labels[idx]
            
            if target_modality is None or modality == target_modality:
                filtered_indices.append(idx)
                filtered_scores.append(float(score))
                filtered_modalities.append(modality)
        
        return filtered_indices, filtered_scores, filtered_modalities
    
    def same_modal_query(
        self,
        query_embedding: torch.Tensor,
        modality: str,
        k: int = 5
    ) -> ModalityResult:
        """
        Query for same modality.
        
        Args:
            query_embedding: Query embedding
            modality: Modality to search
            k: Number of results
            
        Returns:
            ModalityResult with filtered results
        """
        # Search with no filter first
        scores, indices = self.index.search(query_embedding, k=k * 10)  # Get more to filter
        
        # Filter by modality
        filtered_indices, filtered_scores, modalities = self._filter_by_modality(
            indices, scores, target_modality=modality
        )
        
        # Take top-k
        return ModalityResult(
            indices=filtered_indices[:k],
            scores=filtered_scores[:k],
            modalities=modalities[:k],
            query_modality=modality
        )
    
    def cross_modal_query(
        self,
        query_embedding: torch.Tensor,
        source_modality: str,
        target_modality: str,
        k: int = 5
    ) -> ModalityResult:
        """
        Query across modalities.
        
        Args:
            query_embedding: Query embedding
            source_modality: Modality of query image
            target_modality: Modality to search in
            k: Number of results
            
        Returns:
            ModalityResult with filtered results
        """
        # Search with no filter first
        scores, indices = self.index.search(query_embedding, k=k * 10)
        
        # Filter by target modality (excluding source)
        filtered_indices, filtered_scores, modalities = self._filter_by_modality(
            indices, scores, target_modality=target_modality
        )
        
        # Take top-k
        return ModalityResult(
            indices=filtered_indices[:k],
            scores=filtered_scores[:k],
            modalities=modalities[:k],
            query_modality=source_modality
        )
    
    def mixed_query(
        self,
        query_embedding: torch.Tensor,
        source_modality: str,
        k: int = 5
    ) -> ModalityResult:
        """
        Query across all modalities.
        
        Args:
            query_embedding: Query embedding
            source_modality: Modality of query image
            k: Number of results
            
        Returns:
            ModalityResult with results from all modalities
        """
        # Search
        scores, indices = self.index.search(query_embedding, k=k)
        
        # Get modalities
        modalities = [
            self.modality_labels[idx]
            for idx in indices[0]
            if idx >= 0
        ]
        
        return ModalityResult(
            indices=indices[0].tolist(),
            scores=scores[0].tolist(),
            modalities=modalities,
            query_modality=source_modality
        )
    
    def get_modality_distribution(self) -> Dict[str, int]:
        """
        Get distribution of modalities in index.
        
        Returns:
            Dict mapping modality to count
        """
        dist = {}
        for mod in self.modality_labels:
            dist[mod] = dist.get(mod, 0) + 1
        return dist


# Self-check
if __name__ == "__main__":
    print("Testing MultiModalRetrieval...")
    
    # Create dummy embeddings
    n_per_modality = 50
    embed_dim = 768
    
    embeddings_by_modality = {
        "optical": torch.randn(n_per_modality, embed_dim),
        "sar": torch.randn(n_per_modality, embed_dim),
        "multispectral": torch.randn(n_per_modality, embed_dim),
    }
    
    # Normalize
    for mod in embeddings_by_modality:
        embeddings_by_modality[mod] = torch.nn.functional.normalize(
            embeddings_by_modality[mod], dim=1
        )
    
    # Build index
    retrieval = MultiModalRetrieval(embed_dim)
    retrieval.build_index(embeddings_by_modality)
    
    print(f"Index size: {retrieval.size}")
    print(f"Modality distribution: {retrieval.get_modality_distribution()}")
    
    # Same-modal query
    query = torch.randn(embed_dim)
    query = torch.nn.functional.normalize(query, dim=0)
    
    result = retrieval.same_modal_query(query, modality="optical", k=5)
    print(f"\nSame-modal (optical→optical):")
    print(f"  Results: {len(result.indices)}")
    print(f"  Modalities: {result.modalities}")
    
    # Cross-modal query
    result = retrieval.cross_modal_query(
        query, source_modality="optical", target_modality="sar", k=5
    )
    print(f"\nCross-modal (optical→sar):")
    print(f"  Results: {len(result.indices)}")
    print(f"  Modalities: {result.modalities}")
    
    print("\nMultiModalRetrieval test passed!")