import unittest
import os
import shutil
import numpy as np
import sys

# Ensure we can import from src
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))

from sage_lib.partition.PartitionManager import PartitionManager
from sage_lib.single_run.SingleRun import SingleRun

class TestPartitionFAISS(unittest.TestCase):
    def setUp(self):
        self.test_dir = os.path.abspath("./test_partition_faiss_db")
        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_partition_manager_faiss_delegation(self):
        print("\nTesting PartitionManager FAISS basic delegation...")
        # 1. Initialize with HNSW to test efSearch
        pm = PartitionManager(
            path=self.test_dir,
            storage='hybrid',
            use_faiss=True,
            faiss_dim=4,
            faiss_factory_str="HNSW32",
            faiss_metric="L2"
        )
        
        # 2. Add some data
        for i in range(5):
            sr = SingleRun(name=f"run_{i}")
            # MUST set on AtomPositionManager for HybridStorage to extract it for SQL index
            sr.AtomPositionManager.energy = float(i * 10) 
            
            # Add to store
            oid = pm.add(sr)
            # Add vectors manually (since we didn't provide faiss_vec_fn)
            vec = np.array([i, i, i, i], dtype=np.float32)
            pm._store.add_vector(oid, vec)
            
        self.assertEqual(len(pm), 5)
        
        # 3. Test search_vector
        query_vec = np.array([2.1, 2.1, 2.1, 2.1], dtype=np.float32)
        results = pm.search_vector(query_vec, threshold=1.0, k=1)
        self.assertEqual(len(results), 1)
        self.assertEqual(pm.get_container(results[0]).name, "run_2")
        
        # 4. Test search_with_filter
        # Filter where E < 25 (runs 0, 1, 2)
        results = pm.search_with_filter(query_vec, "energy < 25", k=5)
        names = [pm.get_container(r).name for r in results]
        self.assertTrue(len(names) >= 3, f"Expected at least 3 results, got {names}")
        
        # 5. Test novelty scores
        new_vecs = np.array([[10, 10, 10, 10]], dtype=np.float32)
        scores = pm.get_novelty_scores(new_vecs, k=1)
        self.assertTrue(scores[0] > 5.0) 
        
        # 6. Test diversity picking
        diverse_ids = pm.pick_diverse_subset(count=2)
        self.assertEqual(len(diverse_ids), 2)
        
        # 7. Test HNSW tuning
        pm.faiss_efSearch = 64
        self.assertEqual(pm.faiss_efSearch, 64)
        print("HNSW efSearch OK")
        
        pm.flush_faiss()
        pm._store.close()
        
    def test_partition_manager_composite_faiss(self):
        base_dir = os.path.join(self.test_dir, "base")
        local_dir = os.path.join(self.test_dir, "local")
        
        # Create base
        pm_base = PartitionManager(path=base_dir, storage='hybrid', use_faiss=True, faiss_dim=4)
        oid_base = pm_base.add(SingleRun(name="base_1"))
        pm_base._store.add_vector(oid_base, np.array([1,1,1,1], dtype=np.float32))
        pm_base.flush_faiss()
        pm_base._store.close()
        
        # Correct path
        actual_base_root = os.path.join(base_dir, "hybrid_local")
        
        # Create composite
        pm = PartitionManager(
            path=self.test_dir,
            storage='composite',
            base_root=actual_base_root,
            local_root=local_dir, # will create hybrid_local inside this dir
            use_faiss=True,
            faiss_dim=4
        )
        
        # Add to local
        sr_local = SingleRun(name="local_1")
        comp_id_local = pm.add(sr_local)
        
        # We need to add the vector to the local backend. 
        part, bid = pm._store._resolve(comp_id_local)
        self.assertEqual(part, "local")
        pm._store.local.add_vector(bid, np.array([2,2,2,2], dtype=np.float32))
        
        # Search effectively (threshold=10.0 TO INCLUDE BOTH)
        query = np.array([1.1, 1.1, 1.1, 1.1], dtype=np.float32)
        results = pm.search_vector(query, threshold=10.0, k=5)
        
        self.assertTrue(len(results) >= 2, f"Expected 2 results, got {results}")
        # Ordered by distance: [1,1,1,1] is closer than [2,2,2,2]
        self.assertEqual(pm.get_container(results[0]).name, "base_1")
        self.assertEqual(pm.get_container(results[1]).name, "local_1")
        
        pm._store.close()

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