rpa_vision_v3/tests/unit/test_faiss_ivf_optimization.py

"""
Tests pour l'optimisation FAISS IVF

Valide:
- Migration automatique Flat → IVF
- Entraînement automatique de l'index IVF
- Optimisation périodique
- Calcul de nlist optimal
"""

import pytest
import numpy as np
from pathlib import Path
import tempfile
import shutil

from core.embedding.faiss_manager import FAISSManager


class TestFAISSIVFOptimization:
    """Tests pour l'optimisation IVF"""
    
    def setup_method(self):
        """Setup avant chaque test"""
        self.dimensions = 512
        self.temp_dir = Path(tempfile.mkdtemp())
    
    def teardown_method(self):
        """Cleanup après chaque test"""
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
    
    def test_ivf_training(self):
        """Test entraînement automatique de l'index IVF"""
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            auto_optimize=False
        )
        
        # Au début, pas entraîné
        assert not manager.is_trained
        
        # Ajouter des vecteurs (moins de 100)
        for i in range(50):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Toujours pas entraîné
        assert not manager.is_trained
        
        # Ajouter plus de vecteurs pour déclencher l'entraînement
        for i in range(50, 150):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Maintenant entraîné
        assert manager.is_trained
        assert manager.index.ntotal == 150
    
    def test_nlist_calculation(self):
        """Test calcul de nlist optimal"""
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="Flat",
            auto_optimize=False
        )
        
        # Test différentes tailles
        assert manager._calculate_nlist(100) == 100  # min
        assert manager._calculate_nlist(10000) == 100  # sqrt(10000) = 100
        assert manager._calculate_nlist(40000) == 200  # sqrt(40000) = 200
        assert manager._calculate_nlist(1000000) == 1000  # sqrt(1000000) = 1000
    
    def test_auto_migration_flat_to_ivf(self):
        """Test migration automatique Flat → IVF"""
        # Créer index Flat avec seuil bas pour test
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="Flat",
            metric="cosine",
            auto_optimize=True
        )
        
        # Réduire le seuil pour test
        manager.migration_threshold = 100
        
        # Vérifier qu'on commence avec Flat
        assert manager.index_type == "Flat"
        
        # Ajouter des vecteurs jusqu'au seuil
        for i in range(110):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Devrait avoir migré vers IVF
        assert manager.index_type == "IVF"
        assert manager.is_trained
        assert manager.index.ntotal == 110
    
    def test_ivf_search_quality(self):
        """Test qualité de recherche avec IVF"""
        # Créer index IVF
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            nlist=50,
            nprobe=20  # Augmenter nprobe pour meilleure qualité
        )
        
        # Ajouter des vecteurs
        vectors = []
        for i in range(200):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            vectors.append(vector)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Vérifier que l'index contient bien 200 vecteurs
        assert manager.index.ntotal == 200, f"Expected 200 vectors, got {manager.index.ntotal}"
        
        # Rechercher avec un vecteur après l'entraînement (index 150)
        # Les 100 premiers sont utilisés pour l'entraînement
        query = vectors[150]
        results = manager.search_similar(query, k=5)
        
        # Le premier résultat devrait être le vecteur lui-même
        assert len(results) > 0, "No results returned"
        
        # Pour IVF, la recherche est approximative, donc on vérifie juste
        # que le vecteur est dans les résultats (pas forcément en premier)
        embedding_ids = [r.embedding_id for r in results]
        assert "emb_150" in embedding_ids, f"emb_150 not found in results: {embedding_ids}"
        
        # Vérifier qu'au moins un résultat a une bonne similarité
        max_similarity = max(r.similarity for r in results)
        assert max_similarity > 0.8, f"Max similarity too low: {max_similarity}"
    
    def test_ivf_nprobe_effect(self):
        """Test effet de nprobe sur la qualité de recherche"""
        # Créer index avec nprobe faible
        manager_low = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            nlist=50,
            nprobe=1  # Très faible
        )
        
        # Créer index avec nprobe élevé
        manager_high = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            nlist=50,
            nprobe=20  # Élevé
        )
        
        # Ajouter les mêmes vecteurs aux deux
        vectors = []
        for i in range(200):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            vectors.append(vector)
            manager_low.add_embedding(f"emb_{i}", vector)
            manager_high.add_embedding(f"emb_{i}", vector)
        
        # Rechercher avec un vecteur
        query = vectors[0]
        results_low = manager_low.search_similar(query, k=10)
        results_high = manager_high.search_similar(query, k=10)
        
        # nprobe élevé devrait donner de meilleurs résultats
        # (plus de résultats ou meilleure similarité moyenne)
        assert len(results_high) >= len(results_low)
    
    def test_optimize_index(self):
        """Test optimisation périodique de l'index"""
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            nlist=50,  # Petit nlist initial
            auto_optimize=False
        )
        
        # Ajouter beaucoup de vecteurs
        for i in range(500):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Vérifier nlist initial
        initial_nlist = manager.index.nlist
        assert initial_nlist == 50
        
        # Optimiser
        manager.optimize_index()
        
        # nlist devrait avoir changé (optimal pour 500 vecteurs ≈ 22)
        # Mais on garde le nlist actuel car pas assez différent
        # Testons avec un cas plus extrême
        
        # Ajouter encore plus de vecteurs
        for i in range(500, 2000):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Maintenant optimiser devrait changer nlist
        manager.optimize_index()
        
        # Pour 2000 vecteurs, optimal ≈ sqrt(2000) ≈ 45
        # Différence avec 50 n'est pas assez grande (< 50%)
        # Donc nlist ne change pas dans ce cas
        
        # Test avec un nlist vraiment sous-optimal
        manager.index.nlist = 10  # Forcer un nlist très bas
        manager.optimize_index()
        
        # Devrait avoir augmenté
        assert manager.index.nlist > 10
    
    def test_save_load_ivf(self):
        """Test sauvegarde/chargement d'index IVF"""
        # Créer et peupler index IVF
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            nlist=50,
            nprobe=8
        )
        
        # Ajouter des vecteurs
        for i in range(200):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector, metadata={"index": i})
        
        # Sauvegarder
        index_path = self.temp_dir / "test_ivf.index"
        metadata_path = self.temp_dir / "test_ivf.meta"
        manager.save(index_path, metadata_path)
        
        # Charger
        loaded_manager = FAISSManager.load(index_path, metadata_path)
        
        # Vérifier
        assert loaded_manager.dimensions == self.dimensions
        assert loaded_manager.index_type == "IVF"
        assert loaded_manager.metric == "cosine"
        assert loaded_manager.index.ntotal == 200
        assert loaded_manager.is_trained
        assert loaded_manager.index.nlist == 50
        assert loaded_manager.index.nprobe == 8
        
        # Vérifier métadonnées
        assert len(loaded_manager.metadata_store) == 200
        assert loaded_manager.metadata_store[0]["metadata"]["index"] == 0
    
    def test_stats_with_ivf(self):
        """Test statistiques avec index IVF"""
        manager = FAISSManager(
            dimensions=self.dimensions,
            index_type="IVF",
            metric="cosine",
            nlist=50
        )
        
        # Ajouter des vecteurs
        for i in range(200):
            vector = np.random.randn(self.dimensions).astype(np.float32)
            manager.add_embedding(f"emb_{i}", vector)
        
        # Obtenir stats
        stats = manager.get_stats()
        
        # Vérifier
        assert stats["index_type"] == "IVF"
        assert stats["total_vectors"] == 200
        assert stats["is_trained"] == True
        assert stats["nlist"] == 50
        assert stats["nprobe"] == 8
        assert "optimal_nlist" in stats
        assert "nlist_efficiency" in stats


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