src/retrieval/engine.py

"""
Retrieval engine combining feature extraction and FAISS search.

Provides high-level API for image retrieval.
"""

import torch
import time
from PIL import Image
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass

from ..features.extractor import FeatureExtractor
from .index import FAISSIndex


@dataclass
class RetrievalResult:
    """Result of a single retrieval query."""
    indices: List[int]
    scores: List[float]
    query_time_ms: float
    modality: str


class RetrievalEngine:
    """
    High-level retrieval engine.
    
    Combines feature extraction with FAISS search for fast image retrieval.
    """
    
    def __init__(
        self,
        feature_extractor: Optional[FeatureExtractor] = None,
        index: Optional[FAISSIndex] = None,
        device: Optional[str] = None
    ):
        """
        Initialize retrieval engine.
        
        Args:
            feature_extractor: Feature extractor (creates new if None)
            index: FAISS index (creates new if None)
            device: Device to use
        """
        self.feature_extractor = feature_extractor or FeatureExtractor(device=device)
        self.index = index or FAISSIndex(embed_dim=self.feature_extractor.embed_dim)
        
        # Timing statistics
        self._query_times: List[float] = []
    
    def build_index(
        self,
        embeddings: torch.Tensor,
        save_path: Optional[str] = None
    ) -> None:
        """
        Build index from pre-computed embeddings.
        
        Args:
            embeddings: Gallery embeddings, shape (N, embed_dim)
            save_path: Optional path to save index
        """
        self.index.build(embeddings)
        
        if save_path:
            self.index.save(save_path)
    
    def query(
        self,
        image: Image.Image,
        modality: str = "optical",
        k: int = 5
    ) -> RetrievalResult:
        """
        Query with a single image.
        
        Args:
            image: Query image
            modality: Image modality
            k: Number of results
            
        Returns:
            RetrievalResult with indices, scores, and timing
        """
        start_time = time.perf_counter()
        
        # Extract features
        query_embedding = self.feature_extractor.extract_features(
            image, modality=modality, normalize=True
        )
        
        # Search
        scores, indices = self.index.search(query_embedding, k=k)
        
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        self._query_times.append(elapsed_ms)
        
        return RetrievalResult(
            indices=indices[0].tolist(),
            scores=scores[0].tolist(),
            query_time_ms=elapsed_ms,
            modality=modality
        )
    
    def batch_query(
        self,
        images: List[Image.Image],
        modality: str = "optical",
        k: int = 5
    ) -> List[RetrievalResult]:
        """
        Query with multiple images.
        
        Args:
            images: List of query images
            modality: Image modality
            k: Number of results
            
        Returns:
            List of RetrievalResult
        """
        results = []
        
        for image in images:
            result = self.query(image, modality=modality, k=k)
            results.append(result)
        
        return results
    
    def get_timing_stats(self) -> Dict[str, float]:
        """
        Get timing statistics.
        
        Returns:
            Dict with mean, median, p95, p99 query times
        """
        if not self._query_times:
            return {"mean": 0, "median": 0, "p95": 0, "p99": 0}
        
        times = sorted(self._query_times)
        n = len(times)
        
        return {
            "mean": sum(times) / n,
            "median": times[n // 2],
            "p95": times[int(n * 0.95)] if n >= 20 else times[-1],
            "p99": times[int(n * 0.99)] if n >= 100 else times[-1],
        }
    
    @property
    def _query_times(self) -> List[float]:
        """Query times list (lazy init)."""
        if not hasattr(self, '_query_times_list'):
            self._query_times_list = []
        return self._query_times_list


# Self-check
if __name__ == "__main__":
    print("Testing RetrievalEngine...")
    
    # Create dummy data
    n_gallery = 50
    embed_dim = 768
    
    # Build index
    embeddings = torch.randn(n_gallery, embed_dim)
    embeddings = torch.nn.functional.normalize(embeddings, dim=1)
    
    # Initialize engine (without model for testing)
    engine = RetrievalEngine.__new__(RetrievalEngine)
    engine.index = FAISSIndex(embed_dim)
    engine._query_times_list = []
    
    # Build index
    engine.build_index(embeddings)
    print(f"Index built with {engine.index.size} embeddings")
    
    # Simulate query timing
    for _ in range(10):
        start = time.perf_counter()
        query = torch.randn(embed_dim)
        query = torch.nn.functional.normalize(query, dim=0)
        scores, indices = engine.index.search(query, k=5)
        elapsed = (time.perf_counter() - start) * 1000
        engine._query_times_list.append(elapsed)
    
    # Get stats
    stats = engine.get_timing_stats()
    print(f"Timing stats: {stats}")
    
    print("\nRetrievalEngine test passed!")