import os
import shutil
import numpy as np
import sys
import time
import matplotlib.pyplot as plt

# 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):
        self.metadata = {}
        self.atomLabelsList = ['H', 'O']

class MockObj:
    def __init__(self, eid):
        self.E = float(eid)
        self.hash = str(eid)
        self.AtomPositionManager = DummyAPM()

def generate_clustered_data(n_total, dim, n_clusters, cluster_std=0.05):
    """
    Generates a synthetic dataset with clusters to make the threshold sweep interesting.
    """
    print(f"Generating {n_total} vectors in {n_clusters} clusters...")
    n_per_cluster = n_total // n_clusters
    centers = np.random.uniform(0, 1, (n_clusters, dim)).astype('float32')
    
    data = []
    for center in centers:
        cluster = center + np.random.normal(0, cluster_std, (n_per_cluster, dim)).astype('float32')
        data.append(cluster)
    
    data = np.vstack(data)
    # Shuffle
    np.random.shuffle(data)
    return data[:n_total]

def run_threshold_sweep_benchmark():
    print("\n=== FAISS Threshold Sweep Benchmark ===")
    store_dir = "/Users/dimitry/Documents/Code/SAGE/test_sweep_store"
    if os.path.exists(store_dir):
        shutil.rmtree(store_dir)

    dim = 128
    n_structures = 100000 
    n_clusters = 500
    
    vectors = generate_clustered_data(n_structures, dim, n_clusters, cluster_std=0.01)
    
    # Use IVF index for 1M structures to keep it fast
    # IVF1024,Flat means 1024 voronoi cells.
    # Metric L2
    store = HybridStorage(store_dir, use_faiss=True, faiss_dim=dim, faiss_factory_str="IVF1024,Flat")
    
    # IVF index needs training
    print("Training FAISS IVF index...")
    store.faiss_index.index.train(vectors[:10000]) # train with a sample

    print(f"Adding {n_structures} structures and vectors to store (batch mode)...")
    t_add_0 = time.time()
    # Adding 1M objects in chunks to avoid memory issues with list of MockObj
    chunk_size = 100000
    for i in range(0, n_structures, chunk_size):
        idx_end = min(i + chunk_size, n_structures)
        objs = [MockObj(j) for j in range(i, idx_end)]
        ids = store.add_many(objs)
        store.add_many_vectors(ids, vectors[i:idx_end])
        print(f"  Progress: {idx_end}/{n_structures} added.")
    
    t_add_1 = time.time()
    print(f"Total insertion ({n_structures} vectors) took: {t_add_1 - t_add_0:.2f}s")

    thresholds = np.linspace(0.01, 1.0, 50)
    
    # 1. Slow Method (Range Search) - ONLY if N < 15,000
    slow_isolated = {}
    slow_dedup = {}
    if n_structures <= 15000:
        print("\nRunning Slow Method (Explicit Range Search)...")
        start_slow = time.time()
        # We only sample a few points for slow method if it's too many
        sample_thresholds = thresholds[::10] 
        for t in sample_thresholds:
            slow_isolated[t] = store.find_isolated_count_range_search(t)
            slow_dedup[t] = store.find_unique_count_range_search(t)
        end_slow = time.time()
        print(f"Slow method took {end_slow - start_slow:.2f}s for {len(sample_thresholds)} thresholds.")
    else:
        print("\nSkipping Slow Method (O(N^2)) for large dataset.")

    # 2. Fast Method (NN sorting trick)
    print("\nRunning Fast Method (NN-Sorting Trick)...")
    start_fast = time.time()
    fast_results = store.get_threshold_sweep_data(thresholds)
    end_fast = time.time()
    print(f"Fast method took {end_fast - start_fast:.2f}s for {len(thresholds)} thresholds.")

    # Comparison check
    if slow_isolated:
        print("\nVerification (Fast Isolated vs Slow Isolated):")
        for t in sample_thresholds:
            s_val = slow_isolated[t]
            f_val = fast_results[float(t)]
            diff = abs(s_val - f_val)
            status = "OK" if diff == 0 else f"FAILED (diff: {diff})"
            print(f"  Threshold {t:.3f}: Slow={s_val}, Fast={f_val} -> {status}")

    # 3. Plotting
    print("\nGenerating plot: faiss_threshold_sweep.png")
    plt.figure(figsize=(10, 6))
    plt.plot(list(fast_results.keys()), list(fast_results.values()), 'o-', label='Isolated Count (Fast Method)')
    if slow_dedup:
        plt.plot(list(slow_dedup.keys()), list(slow_dedup.values()), 's--', label='Deduplicated Count (Slow Method)')
    
    plt.title(f"FAISS Threshold Sweep ({n_structures} structures, {dim} dimensions)")
    plt.xlabel("Similarity Threshold (L2 Distance Squared)")
    plt.ylabel("Number of Unique/Isolated Structures")
    plt.grid(True, alpha=0.3)
    plt.legend()

    
    plot_path = os.path.join("/Users/dimitry/Documents/Code/SAGE/test", "faiss_threshold_sweep.png")
    plt.savefig(plot_path)
    print(f"Plot saved to: {plot_path}")

    store.close()
    if os.path.exists(store_dir):
        shutil.rmtree(store_dir)

if __name__ == "__main__":
    run_threshold_sweep_benchmark()
