import numpy as np
from sage_lib.miscellaneous.marching_cubes import generate_offset_mesh

def test_periodic_mesh():
    # 1. Define a simple cubic lattice
    lattice = np.eye(3) * 10.0
    pbc = (True, True, True)
    
    # 2. Add an atom at the origin (0,0,0) and another at (8,8,8)
    # With PBC, (0,0,0) is equivalent to (10,10,10)
    points = np.array([
        [0.0, 0.0, 0.0],
        [8.0, 8.0, 8.0]
    ])
    
    # 3. Generate offset mesh at d=2.0
    # At d=2.0, the sphere around (0,0,0) should be partially in the cell
    # and partially wrapped from (10,10,10).
    d = 2.0
    resolution = 30
    verts, faces, normals = generate_offset_mesh(points, d, resolution=resolution, lattice=lattice, pbc=pbc)
    
    print(f"Generated mesh with {len(verts)} vertices and {len(faces)} faces.")
    
    if len(verts) > 0:
        print("Test Passed: Vertices generated.")
        # Check if any vertex is near the periodic boundaries
        on_boundary_x = np.any(np.isclose(verts[:, 0], 0, atol=0.5)) or np.any(np.isclose(verts[:, 0], 10, atol=0.5))
        on_boundary_y = np.any(np.isclose(verts[:, 1], 0, atol=0.5)) or np.any(np.isclose(verts[:, 1], 10, atol=0.5))
        on_boundary_z = np.any(np.isclose(verts[:, 2], 0, atol=0.5)) or np.any(np.isclose(verts[:, 2], 10, atol=0.5))
        
        print(f"Vertices on X boundary: {on_boundary_x}")
        print(f"Vertices on Y boundary: {on_boundary_y}")
        print(f"Vertices on Z boundary: {on_boundary_z}")
    else:
        print("Test Failed: No vertices generated.")

if __name__ == "__main__":
    test_periodic_mesh()
