import numpy as np
import os
import sys

# Add src to path
SAGE_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(os.path.join(SAGE_ROOT, 'src'))

from sage_lib.IO.structure_handling_tools.AtomPosition import AtomPosition

def test_vacuum_box():
    # Case 1: Orthorhombic, centered atoms
    # Lattice 10x10x10
    # atoms spanning X and Y, but centered in Z (2 to 8)
    ap = AtomPosition()
    ap.latticeVectors = np.eye(3) * 10.0
    # Block X and Y by placing atoms across the plane
    for x in [0, 5, 9]:
        for y in [0, 5, 9]:
            ap.add_atom(['H'], [[x, y, 2.0]])
            ap.add_atom(['H'], [[x, y, 8.0]])
    
    # Z fractional: 0.2 and 0.8
    # Gaps: internal 0.6 (len 6), periodic 0.4 (len 4)
    # New code should find the 6.0 gap.
    
    v_box, v_origin = ap.get_vacuum_box(tolerance=0.0)
    print("Case 1: Orthorhombic, centered (internal gap 6.0)")
    print(f"  Vacuum Box Z length: {np.linalg.norm(v_box[2])}")
    # Internal gap starts at 0.2. Origin = 0.2 * 10 = 2.0.
    print(f"  Vacuum Origin Z:     {v_origin[2]}")
    
    # Case 2: Unit mismatch check
    # tolerance = 1.0
    # Origin should be 2.0 + 1.0 = 3.0
    v_box_t, v_origin_t = ap.get_vacuum_box(tolerance=1.0)
    print("\nCase 2: Tolerance check (tolerance=1.0)")
    print(f"  Vacuum Origin Z (expected 3.0): {v_origin_t[2]}")
    print(f"  Vacuum Box Z length (expected 6.0 - 2.0 = 4.0): {np.linalg.norm(v_box_t[2])}")

    # Case 3: Atoms at boundaries (0.1 and 0.9)
    # Gap in middle is from 0.1 to 0.9 (len 8)
    # Periodic gap is from 0.9 to 1.1 (len 2)
    ap2 = AtomPosition()
    ap2.latticeVectors = np.eye(3) * 10.0
    for x in [0, 5, 9]:
        for y in [0, 5, 9]:
            ap2.add_atom(['H'], [[x, y, 1.0]])
            ap2.add_atom(['H'], [[x, y, 9.0]])
            
    v_box2, v_origin2 = ap2.get_vacuum_box(tolerance=0.0)
    print("\nCase 3: Atoms at boundaries (internal gap 8.0)")
    print(f"  Actual Vacuum Box Z length: {np.linalg.norm(v_box2[2])}")
    print(f"  Vacuum Origin Z (expected 1.0): {v_origin2[2]}")

if __name__ == "__main__":
    test_vacuum_box()
