import os
import sys
import numpy as np
from scipy.spatial import cKDTree

# 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.partition.partition_builder.MoleculeCluster_builder import MoleculeCluster_builder
from sage_lib.IO.structure_handling_tools.AtomPositionManager import AtomPositionManager
from sage_lib.IO.structure_handling_tools.AtomPosition import AtomPosition

def check_overlaps(apm, tolerance, atoms_per_mol=2):
    pos = apm.atomPositions
    if len(pos) < 2:
        return []
    
    # Use PeriodicCKDTree for verification if lattice exists
    from sage_lib.miscellaneous.periodic_kdtree import PeriodicCKDTree
    L = getattr(apm, "latticeVectors", None)
    if L is not None:
        tree = PeriodicCKDTree(L, pos)
    else:
        tree = cKDTree(pos)
        
    hits = tree.query_ball_tree(tree, r=tolerance)
    overlaps = []
    for i, h in enumerate(hits):
        mol_i = i // atoms_per_mol
        for j in h:
            if i < j:
                mol_j = j // atoms_per_mol
                if mol_i != mol_j:
                    d = np.linalg.norm(pos[i] - pos[j])
                    # Note: we should strictly check minimum image distance, 
                    # but if they are overlaps, they will show up in the ball query.
                    overlaps.append((i, j, d))
    return overlaps

def test_triclinic_overlaps():
    # Setup a non-orthogonal (triclinic) cell
    apm = AtomPositionManager()
    apm.latticeVectors = np.array([
        [10.0, 0.0, 0.0],
        [2.0, 10.0, 0.0],
        [1.0, 1.0, 10.0]
    ])
    
    class MockContainer:
        def __init__(self, apm):
            self.AtomPositionManager = apm
            self.file_location = "./test"
            
    container = MockContainer(apm)
    builder = MoleculeCluster_builder(name="TriclinicTest", file_location=".")
    
    # Define CO template
    co = AtomPosition()
    co.atomLabelsList = np.array(['C', 'O'], dtype=object)
    co.atomPositions = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.15]])
    builder.add_molecule_template("CO", co)
    
    tolerance = 1.6
    print("Starting batch placement in triclinic cell...")
    # Place enough molecules to force boundary interactions
    builder._add_molecules_batch(
        container=container,
        m_name="CO",
        m_count=100,
        shape='box',
        cluster_lattice_vectors=apm.latticeVectors,
        tolerance=tolerance,
        batch_size=32,
        verbosity=True
    )
    
    overlaps = check_overlaps(apm, tolerance - 1e-5)
    if overlaps:
        print(f"FAILED: Found {len(overlaps)} overlaps in triclinic cell!")
        for i, j, d in overlaps[:5]:
            print(f"  Atoms {i} and {j} at distance {d:.4f} < {tolerance}")
    else:
        print("SUCCESS: No overlaps found in triclinic cell.")

if __name__ == "__main__":
    test_triclinic_overlaps()
