src/retrieval/index.py

"""
FAISS index for fast similarity search.

Handles index building, searching, and persistence.
"""

import faiss
import torch
import numpy as np
from pathlib import Path
from typing import Tuple, Optional


class FAISSIndex:
    """
    FAISS index for cosine similarity search.
    
    Uses IndexFlatIP (inner product) which works as cosine similarity
    when embeddings are L2-normalized.
    
    Supports:
    - Global embeddings (standard CLIP)
    - Multiscale embeddings (fused global + patch features)
    - Multiple index types for different use cases
    """
    
    def __init__(self, embed_dim: int = 768, index_type: str = "flat"):
        """
        Initialize FAISS index.
        
        Args:
            embed_dim: Embedding dimension
            index_type: Type of index ("flat", "ivf", "pq")
        """
        self.embed_dim = embed_dim
        self.index_type = index_type
        self.is_built = False
        self._n_embeddings = 0
        
        # Create index based on type
        if index_type == "flat":
            self.index = faiss.IndexFlatIP(embed_dim)
        elif index_type == "ivf":
            # IVF index for faster search on large datasets
            quantizer = faiss.IndexFlatIP(embed_dim)
            self.index = faiss.IndexIVFFlat(quantizer, embed_dim, 100)
        elif index_type == "pq":
            # Product quantization for memory efficiency
            self.index = faiss.IndexPQ(embed_dim, 8, 8)
        else:
            raise ValueError(f"Unknown index type: {index_type}")
    
    @property
    def size(self) -> int:
        """Number of embeddings in index."""
        return self._n_embeddings
    
    def build(self, embeddings: torch.Tensor, train_index: bool = False) -> None:
        """
        Build index from embeddings.
        
        Args:
            embeddings: Tensor of shape (N, embed_dim), L2-normalized
            train_index: Whether to train IVF/PQ index (requires enough data)
        """
        if isinstance(embeddings, torch.Tensor):
            embeddings = embeddings.numpy().astype(np.float32)
        
        if embeddings.ndim == 1:
            embeddings = embeddings.reshape(1, -1)
        
        assert embeddings.shape[1] == self.embed_dim, \
            f"Expected dim {self.embed_dim}, got {embeddings.shape[1]}"
        
        # Recreate index with correct type
        if self.index_type == "flat":
            self.index = faiss.IndexFlatIP(self.embed_dim)
        elif self.index_type == "ivf":
            quantizer = faiss.IndexFlatIP(self.embed_dim)
            self.index = faiss.IndexIVFFlat(quantizer, self.embed_dim, 100)
            if train_index and embeddings.shape[0] >= 100:
                self.index.train(embeddings)
        elif self.index_type == "pq":
            self.index = faiss.IndexPQ(self.embed_dim, 8, 8)
            if train_index and embeddings.shape[0] >= 100:
                self.index.train(embeddings)
        
        self.index.add(embeddings)
        self._n_embeddings = self.index.ntotal
        self.is_built = True
    
    def search(
        self,
        query: torch.Tensor,
        k: int = 5
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for top-K similar embeddings.
        
        Args:
            query: Query embedding(s), shape (embed_dim,) or (N, embed_dim)
            k: Number of results to return
            
        Returns:
            (scores, indices) where:
            - scores: shape (N, k) with similarity scores
            - indices: shape (N, k) with indices into gallery
        """
        if not self.is_built:
            raise RuntimeError("Index not built. Call build() first.")
        
        # Convert to numpy
        if isinstance(query, torch.Tensor):
            query = query.numpy().astype(np.float32)
        
        # Ensure 2D
        if query.ndim == 1:
            query = query.reshape(1, -1)
        
        # Clamp k to index size
        k = min(k, self._n_embeddings)
        
        # Search
        scores, indices = self.index.search(query, k)
        
        return scores, indices
    
    def save(self, path: str) -> None:
        """
        Save index to disk.
        
        Args:
            path: Path to save index (without extension)
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        # Save FAISS index
        faiss.write_index(self.index, str(path))
    
    def load(self, path: str) -> None:
        """
        Load index from disk.
        
        Args:
            path: Path to saved index
        """
        self.index = faiss.read_index(str(path))
        self._n_embeddings = self.index.ntotal
        self.is_built = True
    
    def get_embeddings(self) -> np.ndarray:
        """Get all embeddings from index."""
        if not self.is_built:
            return np.array([])
        
        embeddings = np.array([
            self.index.reconstruct(i)
            for i in range(self._n_embeddings)
        ])
        
        return embeddings
    
    def search_multiscale(
        self,
        query: torch.Tensor,
        k: int = 5,
        global_weight: float = 0.7
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search with weighted global + patch features.
        
        Args:
            query: Query embedding (fused global + patch)
            k: Number of results
            global_weight: Weight for global features (0-1)
            
        Returns:
            (scores, indices)
        """
        if not self.is_built:
            raise RuntimeError("Index not built. Call build() first.")
        
        if isinstance(query, torch.Tensor):
            query = query.numpy().astype(np.float32)
        
        if query.ndim == 1:
            query = query.reshape(1, -1)
        
        k = min(k, self._n_embeddings)
        scores, indices = self.index.search(query, k)
        
        return scores, indices


# Self-check
if __name__ == "__main__":
    print("Testing FAISSIndex...")
    
    # Create dummy embeddings
    n_gallery = 100
    embed_dim = 768
    
    embeddings = torch.randn(n_gallery, embed_dim)
    embeddings = torch.nn.functional.normalize(embeddings, dim=1)
    
    # Build index
    index = FAISSIndex(embed_dim)
    index.build(embeddings)
    
    print(f"Index built with {index.size} embeddings")
    
    # Search
    query = torch.randn(embed_dim)
    query = torch.nn.functional.normalize(query, dim=0)
    
    scores, indices = index.search(query, k=5)
    
    print(f"Query results:")
    print(f"  Scores shape: {scores.shape}")
    print(f"  Indices shape: {indices.shape}")
    print(f"  Top-5 scores: {scores[0]}")
    print(f"  Top-5 indices: {indices[0]}")
    
    # Save/load roundtrip
    import tempfile
    with tempfile.TemporaryDirectory() as tmpdir:
        save_path = Path(tmpdir) / "test_index.faiss"
        index.save(save_path)
        
        loaded_index = FAISSIndex(embed_dim)
        loaded_index.load(save_path)
        
        print(f"\nLoaded index size: {loaded_index.size}")
        
        # Verify search results match
        scores2, indices2 = loaded_index.search(query, k=5)
        assert np.allclose(scores, scores2), "Scores mismatch!"
        assert np.array_equal(indices, indices2), "Indices mismatch!"
    
    print("\nFAISSIndex test passed!")