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

# 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:
    """Mock AtomPositionManager for testing purposes."""
    def __init__(self):
        self.metadata = {}
        self.atomLabelsList = ['H', 'O']

class MockObj:
    """Mock object representing an atomistic structure with energy and hash."""
    def __init__(self, e=0.0):
        self.E = e
        self.hash = str(e)
        self.AtomPositionManager = DummyAPM()

def test_faiss():
    """
    Standard unit test for basic FAISS functionality in HybridStorage.
    Tests addition, vector search, duplicate group identification, and counting methods.
    """
    print("\n--- Running Basic FAISS Unit Tests ---")
    store_dir = "/Users/dimitry/Documents/Code/SAGE/test_faiss_store"
    if os.path.exists(store_dir):
        shutil.rmtree(store_dir)
        
    # Initialize store with dimension 3
    store = HybridStorage(store_dir, use_faiss=True, faiss_dim=3)
    
    # Define sample vectors
    vec1 = np.array([1.0, 0.0, 0.0])
    vec2 = np.array([0.0, 1.0, 0.0])
    vec3 = np.array([1.0, 0.0, 0.1])
    
    # Add objects and their associated vectors
    id1 = store.add(MockObj(1.0))
    store.add_vector(id1, vec1)
    
    id2 = store.add(MockObj(2.0))
    store.add_vector(id2, vec2)
    
    id3 = store.add(MockObj(3.0))
    store.add_vector(id3, vec3)

    # Search: vec1 should be close to vec3. 
    # Distance is 0.1 in the z-axis -> Squared Euclidean: 0.1^2 = 0.01.
    matched = store.search_vector(vec1, threshold=0.005)
    assert set(matched) == {id1}
    
    matched2 = store.search_vector(vec1, threshold=0.02)
    assert set(matched2) == {id1, id3}
    
    # Global duplicate group identification
    groups = store.find_duplicate_groups(threshold=0.02)
    assert len(groups) == 1
    assert set(groups[0]) == {id1, id3}
    
    store.close()
    
    # Re-open and verify persistence
    store2 = HybridStorage(store_dir, use_faiss=True)
    assert store2.faiss_index.ntotal == 3
    
    # Test high-level unique vs repeated counting (Option 1: Range search)
    # Threshold 0.02 is enough to catch {id1, id3} as a pair.
    res1 = store2.count_unique_and_repeated(threshold=0.02, M=1)
    assert res1["unique"] == 1 # id2 is unique
    assert res1["repeated_groups_count"] == 1 # {id1, id3} is one group
    assert res1["total_repeated_structures"] == 2 # 2 items in that group
    
    # Test batch search-based counting (Option 2: Batch search)
    res2 = store2.check_repeated_in_batches(threshold=0.02, M=1)
    assert res2["unique"] == 1
    assert res2["repeated_structures"] == 2
    
    # Test removal
    store2.remove([id1])
    assert store2.faiss_index.ntotal == 2
    
    # Test indexing rebuild
    id4 = store2.add(MockObj(4.0))
    # Use E-value as a simple 3D vector [0, 0, E]
    store2.rebuild_faiss_index(lambda obj: np.array([0.0, 0.0, obj.E]))
    assert store2.faiss_index.ntotal == 3
    
    store2.close()
    print("Basic FAISS tests passed successfully.")

def benchmark_1m_structures():
    """
    Benchmark the performance of FAISS searching and counting.
    Uses random vectors to simulate a large-scale database and measures timing.
    """
    print("\n--- Benchmark: Large Scale FAISS Performance ---")
    store_dir = "/Users/dimitry/Documents/Code/SAGE/test_benchmark_1m"
    if os.path.exists(store_dir):
        shutil.rmtree(store_dir)
        
    dim = 128
    store = HybridStorage(store_dir, use_faiss=True, faiss_dim=dim)
    
    num_structures = 10000
    print(f"Generating and adding {num_structures} random vectors (dim={dim})...")
    
    t0 = time.time()
    batch_gen = 1000
    for i in range(0, num_structures, batch_gen):
        n = min(batch_gen, num_structures - i)
        vectors = np.random.random((n, dim)).astype('float32')
        ids = np.arange(i + 1, i + n + 1).astype('int64') # 1-based IDs
        store.faiss_index.add_with_ids(vectors, ids)
        if (i+n) % 20000 == 0:
            print(f"  Progress: {i + n}/{num_structures} vectors indexed.")

    # Inject 100 exact duplicates of the FIRST vector to verify detection
    print("Adding 100 exact duplicates of vector #1 for verification...")
    first_vec = store.faiss_index.index.reconstruct(0).reshape(1, -1)
    dup_vectors = np.repeat(first_vec, 100, axis=0)
    dup_ids = np.arange(num_structures + 1, num_structures + 101).astype('int64')
    store.faiss_index.add_with_ids(dup_vectors, dup_ids)
    
    t1 = time.time()
    print(f"Total insertion ({num_structures + 100} vectors) took: {t1 - t0:.2f}s")
    
    # Run range search-based counting (Option 1: Detailed clustering)
    '''
    print(f"\nRunning count_unique_and_repeated (threshold=1e-5, M=1)...")
    t2a = time.time()
    results_detailed = store.count_unique_and_repeated(threshold=1e-5, M=1)
    t3a = time.time()
    
    print(f"Detailed Results (Option 1):")
    print(f"  - Repeated groups:       {results_detailed['repeated_groups_count']}")
    print(f"  - Repeated structures:   {results_detailed['total_repeated_structures']}")
    print(f"  - Unique structures:     {results_detailed['unique']}")
    print(f"  - Execution time:        {t3a - t2a:.2f}s")
    '''
    # Run batch search-based counting (Option 2: Memory efficient)
    print(f"\nRunning check_repeated_in_batches (threshold=1e-5, M=1)...")
    t2b = time.time()
    results_batch = store.check_repeated_in_batches(threshold=1e-1, M=1, batch_size=100000)
    t3b = time.time()
   
    print(f"Batch Results (Option 2):")
    print(f"  - Repeated structures:   {results_batch['repeated_structures']} (Expected: 101)")
    print(f"  - Unique structures:     {results_batch['unique']}")
    print(f"  - Execution time:        {t3b - t2b:.2f}s")
    
    store.close()
    if os.path.exists(store_dir):
        shutil.rmtree(store_dir)

if __name__ == "__main__":
    # test_faiss() # skip small test
    benchmark_1m_structures()

# ==============================================================================
# Usage Examples / Tutorials
# ==============================================================================

def example_convert_db(db_path: str):
    """How to convert an existing non-FAISS database to a FAISS-indexed one."""
    print(f"\n--- Tutorial: Converting Existing DB ({db_path}) ---")
    store = HybridStorage(db_path, use_faiss=True, faiss_dim=100)
    
    def my_vectorizer(structure):
        """Extracts a mathematical vector from the structure."""
        # For simulation, we create a pseudo-random fingerprint
        np.random.seed(int(structure.E * 1000))
        return np.random.rand(100).astype(np.float32)

    # This iterates through the whole database and builds the .faiss index file
    store.rebuild_faiss_index(my_vectorizer)
    print(f"Done! Database indexed with {store.faiss_index.ntotal} vectors.")
    store.close()

def example_find_duplicate_groups(db_path: str):
    """How to find clusters of near-identical structures across the database."""
    print(f"\n--- Tutorial: Finding Global Clusters in {db_path} ---")
    store = HybridStorage(db_path, use_faiss=True)
    
    # Threshold for IndexFlatL2 is squared distance.
    # If 0.05 is the max distance, threshold = 0.0025
    threshold = 0.0025 
    
    clusters = store.find_duplicate_groups(threshold=threshold)
    print(f"Found {len(clusters)} repeated groups.")
    for i, cluster in enumerate(clusters[:3]): 
        print(f"  Cluster #{i+1} (IDs): {cluster}")
        
    store.close()

def example_query_similar(db_path: str):
    """How to query the database for the Top M most similar structures."""
    print(f"\n--- Tutorial: Querying Top M similar in {db_path} ---")
    store = HybridStorage(db_path, use_faiss=True)
    
    # Suppose we have a new structure and want to find similar ones in the DB
    query_vec = np.random.rand(100).astype(np.float32) 
    M = 5
    # Use infinity to ignore threshold and just get Top K closest
    similar_ids = store.search_vector(query_vec, threshold=float('inf'), k=M)
    
    print(f"Top {M} similar structure IDs: {similar_ids}")
    store.close()
