
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

XYZ_FILE = "test_bulk_hybrid.extxyz"
STORE_PATH = "test_hybrid_bulk_store"

def cleanup():
    if os.path.exists(STORE_PATH):
        shutil.rmtree(STORE_PATH)
    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_hybrid_bulk():
    print("Testing HybridStorage bulk insert from XYZ...")
    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.")

    # Test HybridStorage
    pm = PartitionManager(path=STORE_PATH, storage="hybrid")
    
    start = time.time()
    pm.read_files(XYZ_FILE, verbose=False)
    end = time.time()
    
    print(f"Hybrid: 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_hybrid_bulk()
