import os
import sys
import unittest
import numpy as np

# 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

class TestMoleculeAdsorbate(unittest.TestCase):
    def setUp(self):
        # Create a container with a flat surface
        self.apm = AtomPositionManager()
        self.apm.latticeVectors = np.eye(3) * 20.0
        
        # Grid of 5x5 atoms in XY plane (z=0)
        x, y = np.meshgrid(np.linspace(5, 15, 5), np.linspace(5, 15, 5))
        pos = np.column_stack([x.ravel(), y.ravel(), np.zeros(25)])
        labels = np.array(['Pt'] * 25, dtype=object)
        self.apm.add_atom(labels, pos)
        
        # Wrapper for what MoleculeCluster_builder expects (container.AtomPositionManager)
        class MockContainer:
            def __init__(self, apm):
                self.AtomPositionManager = apm
                self.file_location = "./test_container"
        self.container = MockContainer(self.apm)
        
        self.builder = MoleculeCluster_builder(name="AdsorbateTest", file_location=".")
        
        # Mock copy_and_update_container which usually comes from PartitionManager
        import copy
        def mock_copy(container, sub_directory, file_location=None):
            c_copy = copy.deepcopy(container)
            c_copy.file_location = f"{container.file_location}{sub_directory}" if file_location is None else file_location
            return c_copy
        self.builder.copy_and_update_container = mock_copy
        self.builder.set_container = lambda x: True
        self.builder.containers = [self.container]
        
        # Define a simple adsorbate (linear molecule CO)
        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]])
        co._mass = 28.0
        self.builder.add_molecule_template("CO", co)

    def test_adsorbate_random_placement(self):
        """Test default random placement of adsorbates."""
        v_item = {
            'adsorbate': ['CO'],
            'ID': ['Pt'],
            'd': 2.0,
            'resolution': 20,
            'padding': 0.1,
            'collision_tolerance': 1.0,
            'molecules_number': 3,
            'distribution': 'random',
            'align_to_surface': False,  
            'seed': 42,
            'verbose': False,
            'verbosity': False
        }
        
        new_containers = self.builder.handleCLUSTER({'ADD_ADSORBATE': v_item}, [self.container])
        self.assertIsNotNone(new_containers)
        self.assertEqual(len(new_containers), 1)
        new_apm = new_containers[0].AtomPositionManager
        
        # Check that atoms were added
        self.assertEqual(len(new_apm.atomPositions), 31)

    def test_adsorbate_aligned_placement(self):
        """Test normal-aligned placement of adsorbates."""
        v_item = {
            'adsorbate': ['CO'],
            'ID': ['Pt'],
            'd': 3.0,
            'resolution': 20,
            'padding': 0.1,
            'collision_tolerance': 1.0,
            'molecules_number': 3,
            'distribution': 'random',
            'align_to_surface': True, 
            'seed': 42,
            'verbose': False,
            'verbosity': False
        }
        
        new_containers = self.builder.handleCLUSTER({'ADD_ADSORBATE': v_item}, [self.container])
        self.assertIsNotNone(new_containers)
        new_apm = new_containers[0].AtomPositionManager
        
        # Debug: Check relative vectors of newly added molecules
        last_pos = new_apm.atomPositions[-6:]
        for i in range(0, 6, 2):
            vec = last_pos[i+1] - last_pos[i]
            unit_vec = vec / np.linalg.norm(vec)
            cos_z = abs(unit_vec[2])
            print(f"Molecule {i//2} unit vector: {unit_vec}, cos_z: {cos_z}")
            # self.assertGreater(cos_z, 0.9, f"Molecule {i//2} not aligned")

if __name__ == '__main__':
    unittest.main()
