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

DB_PATH = "test_xyz_bulk.db"
XYZ_FILE = "test_bulk.extxyz"

def cleanup():
    if os.path.exists(DB_PATH):
        os.remove(DB_PATH)
    if os.path.exists("test_xyz_bulk_PM"):
        shutil.rmtree("test_xyz_bulk_PM")
    if os.path.exists(XYZ_FILE):
        os.remove(XYZ_FILE)
    if os.path.exists(XYZ_FILE + ".xzi.npz"):
        os.remove(XYZ_FILE + ".xzi.npz")

def test_xyz_bulk():
    print("Testing XYZ read bulk insert...")
    cleanup()
    
    # create dummy xyz
    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)
    
    write(XYZ_FILE, atoms_list)
    print(f"Generated {XYZ_FILE} with {n_atoms} structures.")

    pm = PartitionManager(path="test_xyz_bulk_PM", storage="sqlite", db_path=DB_PATH)
    
    start = time.time()
    # read_files calls read_XYZ for .extxyz
    pm.read_files(XYZ_FILE, verbose=False)
    end = time.time()
    
    print(f"Read and 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_xyz_bulk()
