import os
import shutil
import unittest
import numpy as np
from sage_lib.partition.PartitionManager import PartitionManager
from sage_lib.single_run.SingleRun import SingleRun

class TestFAISSMultiVectorStress(unittest.TestCase):
    def setUp(self):
        self.test_dir = os.path.abspath("test_stress_faiss")
        if os.path.exists(self.test_dir):
            shutil.rmtree(self.test_dir)
        os.makedirs(self.test_dir)

    def tearDown(self):
        if os.path.exists(self.test_dir):
            shutil.rmtree(self.test_dir)

    def test_multivector_stress(self):
        # 1. Init
        pm = PartitionManager(
            path=self.test_dir,
            storage='hybrid',
            use_faiss=True,
            faiss_dim=8
        )
        
        n_structs = 500 
        n_atoms = 10
        
        print(f"Adding {n_structs} structures with {n_atoms} vectors each...")
        
        # We create some "target" structures that we will search for
        target_indices = [100, 250, 400]
        
        for i in range(n_structs):
            sr = SingleRun(name=f"struct_{i}")
            oid = pm.add(sr)
            
            if i in target_indices:
                # Target vectors: close to [1,1,1,1,1,1,1,1]
                vectors = np.ones((n_atoms, 8), dtype=np.float32) + np.random.normal(0, 0.01, (n_atoms, 8)).astype(np.float32)
            else:
                # Random noise vectors (mean 10, std 1) - far from target
                vectors = (np.random.normal(10, 1, (n_atoms, 8))).astype(np.float32)
            
            pm._store.add_vector(oid, vectors)
            
        pm.flush_faiss()
        
        # 2. Search
        # Query for one of the target vectors
        query_vec = np.ones(8, dtype=np.float32)
        
        print("Searching with unique=True...")
        # threshold=1.0 is plenty for vectors close to [1..1]
        results = pm.search_vector(query_vec, threshold=1.0, k=10, unique=True)
        
        # Verify we found our targets
        found_names = [pm.get_container(r).name for r in results]
        print(f"Found structures: {found_names}")
        
        for idx in target_indices:
            self.assertIn(f"struct_{idx}", found_names)
            
        # Verify unique worked (no duplicates in results)
        self.assertEqual(len(results), len(set(results)))
        
        pm._store.close()

if __name__ == "__main__":
    unittest.main()
