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.miscellaneous.periodic_kdtree import PeriodicCKDTree

def test_periodic_kdtree():
    L = np.eye(3) * 10.0
    # One atom at [0.1, 0.1, 0.1]
    X = np.array([[0.1, 0.1, 0.1]])
    tree = PeriodicCKDTree(L, X)
    
    # Check neighbor near [0.1, 0.1, 0.1]
    q1 = np.array([[0.2, 0.1, 0.1]])
    hits1 = tree.query_ball_point(q1, r=0.2)
    print(f"Direct hits (expected [[0]]): {hits1}")
    
    # Check neighbor across boundary: [9.9, 0.1, 0.1] should hit [0.1, 0.1, 0.1] across X boundary
    q2 = np.array([[9.9, 0.1, 0.1]])
    hits2 = tree.query_ball_point(q2, r=0.3)
    print(f"Boundary hits (expected [[0]]): {hits2}")
    
    # Check batch of points
    q3 = np.array([[0.2, 0.1, 0.1], [9.9, 0.1, 0.1]])
    hits3 = tree.query_ball_point(q3, r=0.3)
    print(f"Batch hits (expected [[0], [0]]): {hits3}")

if __name__ == "__main__":
    try:
        test_periodic_kdtree()
    except Exception as e:
        print(f"Error: {e}")
