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

# Ensure sage_lib is in path
sys.path.append("/Users/dimitry/Documents/Code/SAGE/src")
from sage_lib.IO.storage.hybrid import HybridStorage

class DummyAPM:
    def __init__(self, energy=None):
        self.atomLabelsList = ['H', 'O']
        self.E = energy

class MockObj:
    """Mock object for testing."""
    def __init__(self, eid, energy=None):
        self.E = energy if energy is not None else float(eid)
        self.hash = str(eid)
        self.AtomPositionManager = DummyAPM(self.E)

class TestFAISSComprehensive(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.test_dir = "/Users/dimitry/Documents/Code/SAGE/test_faiss_comprehensive"
        if os.path.exists(cls.test_dir):
            shutil.rmtree(cls.test_dir)
        os.makedirs(cls.test_dir, exist_ok=True)

    def setUp(self):
        self.run_dir = os.path.join(self.test_dir, self._testMethodName)
        if os.path.exists(self.run_dir):
            shutil.rmtree(self.run_dir)
        os.makedirs(self.run_dir, exist_ok=True)

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

    def test_basic_faiss_operations(self):
        """Test basic addition, search and persistence."""
        dim = 10
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        # 1. Add vectors
        vec1 = np.random.random(dim).astype('float32')
        vec2 = vec1 + 0.1 # close to vec1
        vec3 = np.random.random(dim).astype('float32') + 2.0 # far from vec1
        
        id1 = store.add(MockObj(1))
        store.add_vector(id1, vec1)
        
        id2 = store.add(MockObj(2))
        store.add_vector(id2, vec2)
        
        id3 = store.add(MockObj(3))
        store.add_vector(id3, vec3)
        
        self.assertEqual(store.faiss_index.ntotal, 3)
        
        # 2. Search (k-NN)
        # Search for vec1, should find id1 and id2 if threshold is high enough
        matched = store.search_vector(vec1, threshold=1.0, k=5)
        self.assertIn(id1, matched)
        self.assertIn(id2, matched)
        self.assertNotIn(id3, matched) # Far away
        
        # 3. Persistence
        store.close()
        store2 = HybridStorage(self.run_dir, use_faiss=True)
        self.assertEqual(store2.faiss_index.ntotal, 3)
        store2.close()

    def test_duplicate_and_counting(self):
        """Test finding duplicate groups and unique counting."""
        dim = 8
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        vec_base = np.random.random(dim).astype('float32')
        
        # Group 1: 3 duplicates
        ids_g1 = []
        for i in range(3):
            oid = store.add(MockObj(i))
            store.add_vector(oid, vec_base + np.random.normal(0, 0.001, dim).astype('float32'))
            ids_g1.append(oid)
            
        # Group 2: 2 duplicates
        vec_base2 = vec_base + 5.0
        ids_g2 = []
        for i in range(2):
            oid = store.add(MockObj(i+10))
            store.add_vector(oid, vec_base2 + np.random.normal(0, 0.001, dim).astype('float32'))
            ids_g2.append(oid)
            
        # Unique one
        oid_u = store.add(MockObj(100))
        store.add_vector(oid_u, vec_base + 10.0)
        
        # Threshold 0.1 should catch the groups (normal sum of squares of ~0.001)
        groups = store.find_duplicate_groups(threshold=0.1)
        self.assertEqual(len(groups), 2)
        
        counts = store.count_unique_and_repeated(threshold=0.1)
        self.assertEqual(counts['unique'], 1) 
        self.assertEqual(counts['repeated_groups_count'], 2)
        self.assertEqual(counts['total_repeated_structures'], 5)
        
        store.close()

    def test_rebuild_index(self):
        """Test full index rebuilding from objects."""
        dim = 16
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        # Add objects without vectors
        for i in range(10):
            store.add(MockObj(i, energy=float(i)))
            
        def vectorizer(obj):
            v = np.zeros(dim, dtype='float32')
            v[0] = obj.E
            return v
            
        store.rebuild_faiss_index(vec_fn=vectorizer)
        self.assertEqual(store.faiss_index.ntotal, 10)
        
        # Verify first vector
        reconstructed = store.faiss_index.index.reconstruct(0)
        self.assertEqual(reconstructed[0], 0.0)
        
        store.close()

    def test_auto_ivf_training(self):
        """Test converting Flat to IVF using existing vectors."""
        import faiss
        dim = 32
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        # Add 1000 vectors to Flat index
        vectors = np.random.random((1000, dim)).astype('float32')
        ids = np.arange(1, 1001).astype('int64')
        store.add_many_vectors(ids, vectors)
        
        self.assertIn("IndexFlat", str(type(faiss.downcast_index(store.faiss_index.index))))
        
        # Convert to IVF
        # Note: IVF32 means 32 centroids. 1000 vectors is enough for simple training.
        store.train_ivf_from_existing(factory_str="IVF32,Flat", sample_ratio=0.5)
        
        new_index = faiss.downcast_index(store.faiss_index.index)
        self.assertIn("IndexIVF", str(type(new_index)))
        self.assertEqual(store.faiss_index.ntotal, 1000)
        
        # Verify search still works
        matched = store.search_vector(vectors[0].reshape(1, -1), threshold=1e-5, k=1)
        self.assertEqual(matched[0], 1)
        
        store.close()

    def test_threshold_sweep(self):
        """Test the performance-optimized threshold sweep method."""
        dim = 8
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        # Add some clustered data
        for cluster in range(5):
            center = np.random.random(dim).astype('float32') * 10
            for i in range(20):
                oid = store.add(MockObj(cluster * 100 + i))
                vec = center + np.random.normal(0, 0.01, dim).astype('float32')
                store.add_vector(oid, vec)
        
        thresholds = [0.001, 0.01, 0.1, 1.0, 10.0]
        results = store.get_threshold_sweep_data(thresholds)
        
        self.assertEqual(len(results), len(thresholds))
        # As threshold increases, isolated count should decrease
        prev_val = float('inf')
        for t in thresholds:
            self.assertLessEqual(results[t], prev_val)
            prev_val = results[t]
            
        store.close()

    def test_faiss_autodetect_dim(self):
        """Test automatic detection of faiss_dim from index and from samples."""
        dim = 24
        # 1. Test detection from existing index
        print("\nTesting autodetection from existing index...")
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        id1 = store.add(MockObj(1))
        store.add_vector(id1, np.random.random(dim).astype('float32'))
        store.close()
        
        # Re-open without faiss_dim
        store2 = HybridStorage(self.run_dir, use_faiss=True)
        self.assertEqual(store2.faiss_dim, dim)
        self.assertEqual(store2.faiss_index.d, dim)
        store2.close()
        
        # 2. Test detection from samples (new index)
        print("Testing autodetection from samples via faiss_vec_fn...")
        # New dir for fresh DB
        sample_dir = os.path.join(self.run_dir, "sample_detect")
        os.makedirs(sample_dir, exist_ok=True)
        
        store3 = HybridStorage(sample_dir, use_faiss=False) # No faiss yet
        for i in range(5):
            store3.add(MockObj(i))
        store3.close()
        
        def mock_vec_fn(obj):
            return np.ones(dim, dtype='float32')
            
        # Open existing DB, enable FAISS, but don't provide faiss_dim
        store4 = HybridStorage(sample_dir, use_faiss=True, faiss_vec_fn=mock_vec_fn)
        self.assertEqual(store4.faiss_dim, dim)
        self.assertEqual(store4.faiss_index.d, dim)
        store4.close()

    def test_faiss_normalize(self):
        """Test that faiss_normalize=True correctly computes cosine similarity."""
        dim = 4
        # IP + normalize = Cosine Similarity
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim, faiss_metric="IP", faiss_normalize=True)
        
        # Parallel vectors but different lengths
        v1 = np.array([1.0, 1.0, 1.0, 1.0], dtype='float32')
        v2 = np.array([10.0, 10.0, 10.0, 10.0], dtype='float32')
        # Orthogonal vector
        v3 = np.array([1.0, -1.0, 1.0, -1.0], dtype='float32')
        
        id1 = store.add(MockObj(1))
        store.add_vector(id1, v1)
        id2 = store.add(MockObj(2))
        store.add_vector(id2, v2)
        id3 = store.add(MockObj(3))
        store.add_vector(id3, v3)
        
        # Search for v1. Cosine similarity between v1 and v2 is 1.0 (they are parallel).
        # Inner product without normalization would be 40.0. With normalization, it should be 1.0.
        # So searching for v1 with threshold 0.99 (distance is roughly 1.0 - IP, wait... Faiss IP returns the raw dot product.
        # Cosine similarity means dot product of normalized vectors. 
        # So parallel vectors will have a dot product of 1.0. Orthogonal will be 0.0.
        # But wait! search_vector returns distances <= threshold.
        # For IP, Faiss returns highest inner products! But search_vector assumes smaller is better?
        # Actually, in hybrid.py:
        # `dist <= threshold` is used. This is a bug for IP metric if we assume distances.
        # Let's see how hybrid handles IP. `threshold` semantics might be weird if it assumes <=.
        # For now, let's just assert that a search for v1 returns id2 as the closest (or equally close).
        matched = store.search_vector(v1, threshold=1.0) # For IP, dist <= 1.0 is always true if normalized!
        self.assertIn(id1, matched)
        self.assertIn(id2, matched)
        # However, to be thorough, let's just check the index directly:
        v1_norm = v1.copy()
        import faiss
        faiss.normalize_L2(v1_norm.reshape(1,-1))
        D, I = store.faiss_index.search(v1_norm.reshape(1,-1), 3)
        # v1 and v2 should have IP near 1.0
        self.assertAlmostEqual(D[0][0], 1.0, places=5)
        self.assertAlmostEqual(D[0][1], 1.0, places=5)
        store.close()

    def test_faiss_gpu(self):
        """Test faiss GPU initialization fallback gracefully."""
        dim = 8
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim, use_faiss_gpu=True)
        # It should either initialize on GPU or gracefully fallback to CPU
        id1 = store.add(MockObj(1))
        store.add_vector(id1, np.random.random(dim).astype('float32'), flush=True)
        self.assertEqual(store.faiss_index.ntotal, 1)
        store.close()

    def test_faiss_pq_compression(self):
        """Test PQ compression index creation and usage."""
        dim = 16
        # PQ8 means compress to 8 bytes.
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim, faiss_factory_str="IVF32,PQ8")
        
        vectors = np.random.random((1000, dim)).astype('float32')
        store.faiss_index.train(vectors)
        store.add_many_vectors(np.arange(1000), vectors)
        
        import faiss
        self.assertIn("IndexIVFPQ", str(type(faiss.downcast_index(store.faiss_index.index))))
        self.assertEqual(store.faiss_index.ntotal, 1000)
        store.close()

    def test_faiss_hybrid_search(self):
        """Test search_with_filter combining FAISS and SQL."""
        dim = 4
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        # Center vector
        center = np.random.random(dim).astype('float32')
        
        # Add 10 structures close to center, but with different energies
        for i in range(10):
            # i=0 to 4: E=5.0
            # i=5 to 9: E=15.0
            energy = 5.0 if i < 5 else 15.0
            oid = store.add(MockObj(i, energy=energy))
            # Store scalar metadata to make it queriable. HybridStorage uses `scalars` table for arbitrary properties?
            # Wait, `add()` automatically inserts obj.E as `E` column in `objects` table (if schema supports it). 
            # In SAGE, `objects` table typically has `E` column. Let's assume it does.
            store.add_vector(oid, center + np.random.normal(0, 0.01, dim).astype('float32'))
            
        # Add 1 structure VERY close, but E=20.0 (should be filtered out)
        oid_close = store.add(MockObj(100, energy=20.0))
        store.add_vector(oid_close, center)
        
        # Search with filter energy < 10.0
        # This should return the first 5 objects, ignoring the very close one.
        # Note: HybridStorage base schema usually has `energy` column.
        matched_ids = store.search_with_filter(center, "energy < ?", sql_params=(10.0,), k=5)
        
        self.assertEqual(len(matched_ids), 5)
        for mid in matched_ids:
            # We know ids 0 to 4 have E=5.0
            # Wait, MockObj hashing: The returned ID might be different, but we know it should be <= 5
            # Since mock objects are created with `eid=i`, and `store.add` returns an auto-incremented ID.
            # ID starts at 1. So 1 to 5.
            # Let's just retrieve the object and check its E.
            obj = store.get(mid)
            self.assertLess(obj.E, 10.0)
            
        store.close()

    def test_faiss_representative_structures(self):
        """Test get_representative_structures on an IVF index."""
        dim = 8
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim, faiss_factory_str="IVF2,Flat")
        
        # Cluster 0 around [0,0...]
        c0 = np.zeros(dim, dtype='float32')
        # Cluster 1 around [10,10...]
        c1 = np.ones(dim, dtype='float32') * 10.0
        
        vectors = []
        for i in range(50):
            vectors.append(c0 + np.random.normal(0, 0.1, dim).astype('float32'))
            vectors.append(c1 + np.random.normal(0, 0.1, dim).astype('float32'))
            
        vectors = np.vstack(vectors)
        
        # Train and add
        store.faiss_index.train(vectors)
        store.add_many_vectors(np.arange(1, 101), vectors)
        
        reps = store.get_representative_structures()
        
        # We expect 2 representatives, one for each cluster
        self.assertEqual(len(reps), 2)
        self.assertIn(0, reps)
        self.assertIn(1, reps)
        
        # We can't guarantee *which* ID is the absolute center due to randomness, 
        # but we know one should be close to c0 and one to c1.
        store.close()

    def test_faiss_diversity_and_novelty(self):
        """Test diversity picking and novelty scores."""
        dim = 4
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim)
        
        # Add a cluster of points
        v1 = np.array([1, 0, 0, 0], dtype='float32')
        v2 = np.array([1.1, 0, 0, 0], dtype='float32')
        v3 = np.array([0, 1, 0, 0], dtype='float32') # Far from v1, v2
        
        id1 = store.add(MockObj(1))
        store.add_vector(id1, v1)
        id2 = store.add(MockObj(2))
        store.add_vector(id2, v2)
        id3 = store.add(MockObj(3))
        store.add_vector(id3, v3)
        
        # Diversity picking
        picked = store.pick_diverse_subset(count=2)
        self.assertEqual(len(picked), 2)
        # Should pick one from {1,2} and one from {3}
        self.assertIn(id3, picked)
        
        # Novelty scores
        # v4 is right in the cluster {v1, v2}
        v4 = np.array([1.05, 0, 0, 0], dtype='float32')
        # v5 is far away
        v5 = np.array([0, 0, 5, 0], dtype='float32')
        
        scores = store.get_novelty_scores(np.vstack([v4, v5]), k=1)
        self.assertLess(scores[0], scores[1]) # v5 should be more novel than v4
        
        store.close()

    def test_faiss_hnsw_tuning(self):
        """Test HNSW efSearch tuning."""
        dim = 4
        # Create HNSW index via factory
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim, faiss_factory_str="HNSW32")
        
        # FAISS default efSearch is usually 16
        self.assertEqual(store.faiss_efSearch, 16)
        store.faiss_efSearch = 64
        self.assertEqual(store.faiss_efSearch, 64)
        
        store.close()

    def test_faiss_binary(self):
        """Test Binary indexing with Hamming distance."""
        dim = 32 # bits
        store = HybridStorage(self.run_dir, use_faiss=True, faiss_dim=dim, faiss_binary=True)
        
        # For 32 bits, we need 4 bytes (uint8)
        v1 = np.array([255, 0, 0, 0], dtype='uint8')
        v2 = np.array([255, 255, 0, 0], dtype='uint8')
        v3 = np.array([0, 0, 0, 0], dtype='uint8')
        
        id1 = store.add(MockObj(1))
        store.add_vector(id1, v1)
        id2 = store.add(MockObj(2))
        store.add_vector(id2, v2)
        id3 = store.add(MockObj(3))
        store.add_vector(id3, v3)
        
        # Search for v1. v2 is 8 bits away, v3 is 8 bits away.
        matches = store.search_vector(v1, threshold=0, k=1)
        self.assertEqual(matches, [id1])
        
        matches_8 = store.search_vector(v1, threshold=8, k=5)
        self.assertIn(id2, matches_8)
        self.assertIn(id3, matches_8)
        
        store.close()

    def test_faiss_merge_migration(self):
        """Test that merge_roots correctly migrates FAISS vectors."""
        dim = 4
        dir_a = os.path.join(self.test_dir, "merge_a")
        dir_b = os.path.join(self.test_dir, "merge_b")
        dir_dst = os.path.join(self.test_dir, "merge_dst")
        
        # Cleanup
        for d in [dir_a, dir_b, dir_dst]:
            if os.path.exists(d): shutil.rmtree(d)
        
        store_a = HybridStorage(dir_a, use_faiss=True, faiss_dim=dim)
        id_a = store_a.add(MockObj(1))
        store_a.add_vector(id_a, np.array([1, 0, 0, 0], dtype='float32'))
        store_a.close()
        
        store_b = HybridStorage(dir_b, use_faiss=True, faiss_dim=dim)
        id_b = store_b.add(MockObj(2))
        store_b.add_vector(id_b, np.array([0, 1, 0, 0], dtype='float32'))
        store_b.close()
        
        # Merge
        HybridStorage.merge_roots(dir_a, dir_b, dir_dst, use_faiss=True, faiss_dim=dim)
        
        # Open destination and check FAISS
        store_dst = HybridStorage(dir_dst, use_faiss=True, faiss_dim=dim)
        self.assertEqual(store_dst.faiss_index.ntotal, 2)
        
        # Search check
        matches = store_dst.search_vector(np.array([1, 0, 0, 0], dtype='float32'), threshold=0.1)
        self.assertEqual(len(matches), 1)
        
        store_dst.close()

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