
import os
import time
import shutil
import numpy as np
from ase import Atoms
from sage_lib.partition.PartitionManager import PartitionManager

DB_PATH = "test_bulk_performance.db"

def cleanup():
    if os.path.exists(DB_PATH):
        os.remove(DB_PATH)
    if os.path.exists("test_bulk_performance"):
        shutil.rmtree("test_bulk_performance")

def test_sqlite_bulk():
    print("Testing SQLite bulk insert...")
    cleanup()
    
    # create dummy atoms
    n_atoms = 1000
    atoms_list = []
    for i in range(n_atoms):
        a = Atoms('H2O', positions=np.random.rand(3,3), cell=[10,10,10], pbc=True)
        a.info = {'energy': float(i), 'id': i}
        atoms_list.append(a)

    pm = PartitionManager(path="test_bulk_performance_PM", storage="sqlite", db_path=DB_PATH)
    
    start = time.time()
    pm.add_ase(atoms_list)
    end = time.time()
    
    print(f"Added {n_atoms} structures in {end - start:.4f} seconds.")
    
    # Verify count
    assert len(pm) == n_atoms
    print(f"Count verified: {len(pm)}")
    
    cleanup()

if __name__ == "__main__":
    test_sqlite_bulk()
