import os 
import json
import numpy as np
import warnings as warning
from ase import Atoms
from molbar.io.filereader import FileReader
from dscribe.descriptors import SOAP
from molbar.helper.vector import Vector
from scipy.spatial.transform import Rotation 
from itertools import permutations

class VSEPR:

    def classify_nodes_geometry(self, is_adjacent_visible=False, include_all=False) -> None:
        """
        Classifies the geometry of all nodes of the molecule according to VSEPR theory.

        Args:
            is_adjacent_visible (bool, optional): Whether to consider only visible adjacent nodes for the classification. Defaults to False.
            include_all (bool, optional): Whether to classify all nodes in the classification. Otherwise only visible nodes are classified. Defaults to False.
        
        Returns:
            None
        """

        if self.is_2D:

            self.classify_nodes_geometry_from2D()

            return

        # Load ideal SOAPs from file
        with open(os.path.join(os.path.dirname(__file__), "..", "..", "data", 'soap.json')) as json_file:

            self.ideal_soaps = json.load(json_file)

        if include_all:

            nodes = [node for node, data in self.nodes(data=True)]

        else:

            nodes = [node for node, data in self.nodes(data=True) if data.get("visible") == True]

        # Classify all nodes
        vsepr_classes = [self._classify_node_geometry(node, nodes, is_adjacent_visible=is_adjacent_visible)  for node in nodes]

        # Add classification to nodes
        self.add_node_data(attribute="vsepr_classes", new_data=vsepr_classes)

        del self.ideal_soaps

    def _classify_node_geometry(self, central_node: int, initial_visible_nodes: list, is_adjacent_visible=False) -> str:
        """
        Classifies the geometry of one node of the molecule according to VSEPR theory.

        Args:
            node (int): node id
            nodes (list): list of all visible nodes
            is_adjacent_visible (bool, optional): Whether to consider only visible adjacent nodes for the classification. Defaults to False.

        Returns:
            str: VSEPR class of the node.
        """

        # Get adjacent nodes to central node
        adjacent_nodes = self.return_adjacent_nodes(central_node, is_adjacent_visible=is_adjacent_visible)

        self.set_nodes_visible(visible_nodes=[central_node])

        central_cycle_ids = self.return_node_data(attribute="cycle_id")

        # Check if node is in cycle, has influence on VSEPR class, e.g. usally only linear geometry is allowed, but is allowed to be bent if in cycle
        if self.return_node_data(attribute="cycle_id")[0] is not None: 
            
            is_in_cycle = True

        else:

            is_in_cycle = False

        nuclear_charge = self.return_node_data(attribute="atomic_numbers")[0]

        valence = self.return_node_data(attribute="valences")[0]

        self.set_nodes_visible(visible_nodes=[central_node]+adjacent_nodes)

        visible_nodes = self.return_visible_nodes()
        # Get index of central node in visible nodes
        core_index = visible_nodes.index(central_node)

        # Get coordinates of all visible nodes
        coordinates = self.return_node_data(attribute="coordinates")

        elements = self.return_node_data(attribute="elements")

        # Check if one of the adjacent nodes or central node is a metal
        is_metal = any(self.return_node_data(attribute="is_metal"))

        all_in_cycle = all([all([True if id in central_cycle_ids else False for id in ids]) if ids is not None else False for ids in self.return_node_data(attribute="cycle_id")])

        n_atoms = self.return_n_atoms()

        self.set_nodes_visible(visible_nodes=initial_visible_nodes)

        # Rescale coordinates to unit vectors, so that there is no distance dependency
        coordinates = np.array([Vector(coordinates[[core_index,i]]).unit_vector() for i in range(n_atoms)])
        
        # Calculate SOAP for central node
        soap = self._get_soap(coordinates, n_atoms, core_index)

        # Classify node according to VSEPR theory by comparing SOAP to ideal SOAPs
        vsepr_class = self._classifiy_soap(soap, elements[core_index], nuclear_charge, valence, n_atoms, list(elements), is_metal, is_in_cycle, all_in_cycle)

        return vsepr_class

    def classify_nodes_geometry_from2D(self):
        visible_nodes = self.return_visible_nodes()
        elements = self.return_node_data(attribute="elements")
        degrees = self.get_degree()

        vsepr_classes = [ ]

        # for uncharged species:
        for node, element, degree in zip(visible_nodes, elements, degrees):
            self.set_nodes_visible(visible_nodes=[node])

            # Check if node is in cycle, has influence on VSEPR class, e.g. usally only linear geometry is allowed, but is allowed to be bent if in cycle
            if self.return_node_data(attribute="cycle_id")[0] is not None: 
                is_in_cycle = True

            else:
                is_in_cycle = False

            ### 0 binding partners ####

            if degree == 0:
                vsepr_classes.append("dot")
            
            ### 1 bidning partner ###

            elif degree == 1:
                vsepr_classes.append("linear_cn1")

            ### 2 binding partners ###

            # If node is in cycle, it is allowed to be bent, otherwise it is linear, bent contains no important information
            elif (degree == 2) and (is_in_cycle == False):
                vsepr_classes.append("linear_cn2")

            elif (degree == 2) and (is_in_cycle == True):
                vsepr_classes.append("bent")

            ### 3 binding partners ###

            elif (degree == 3) and (element == "P"):
                vsepr_classes.append("trigonal_pyramidal")

            elif (degree == 3) and (element != "P"):
                vsepr_classes.append("trigonal_planar")

            ### 4 binding partners ###

            elif degree == 4:
                vsepr_classes.append("tetrahedral")

            ### 5 binding partners ###

            elif degree == 5:
                vsepr_classes.append("trigonal_bipyramidal")

            ### 6 binding partners ####

            elif degree == 6:
                vsepr_classes.append("octahedral")

            ### more binding partners ####

            else:
                vsepr_classes.append("unknown")
        
        self.set_nodes_visible(visible_nodes = visible_nodes)

        self.add_node_data(attribute="vsepr_classes", new_data=vsepr_classes)


    def _classifiy_soap(self, soap: list, central_element: str, nuclear_charge: int, valence: int, n_atoms: int, all_elements: list, adjacent_metal=False, is_in_cycle=False, all_in_cycle=False):
        """
        Classifies the geometry of one node of the molecule according to VSEPR theory based on the SOAP of the node.

        Args:
            soap (list): SOAP of the node
            central_element (str): element of the central node
            nuclear_charge (int): nuclear charge of the central node
            n_atoms (int): number of nodes in local geometry
            adjacent_metal (bool, optional): Whether one of the adjacent nodes is a metal. Defaults to False.
            is_in_cycle (bool, optional): Whether the node is in a cycle. Defaults to False.
            in_chain (bool, optional): Whether the node is in a chain. Defaults to False.
            all_in_cycle (bool, optional): Whether all nodes (central + adjacent) are in a cycle. Defaults to False.

        Returns:
            str: VSEPR class of the node.
        """
  
        if n_atoms == 1:

            return "dot"

        elif n_atoms == 2:

            return "linear_cn1"
        
        ################## n_adj=3 ##################

        # If node is in cycle, it is allowed to be bent, otherwise it is linear, bent contains no important information
        elif (n_atoms == 3) and (central_element != "N" and central_element != "C" and central_element !="B") and (is_in_cycle == False) and (valence != 2):

            return self._get_best_fit(soap, ["bent", "linear_cn2"])
        
        elif (n_atoms == 3) and (central_element == "B"):

            return "linear_cn2"
        
        elif (n_atoms == 3) and (central_element != "N" and central_element != "C") and (is_in_cycle == False) and (valence == 2):

            return "linear_cn2"
        
        elif (n_atoms == 3) and (central_element == "C") and (is_in_cycle == False) and (valence == 4):

            return "linear_cn2"
        
        elif (n_atoms == 3) and (central_element == "N") and (is_in_cycle == False) and (valence == 4):

            return "bent"
        
        elif (n_atoms == 3) and (central_element == "N") and (is_in_cycle == False) and (valence != 4):

            return "bent"
        
        elif (n_atoms == 3) and (is_in_cycle == True):

            return "bent"
        
        ################## n_adj=4 ##################

        # If element is not C, N or O, it is allowed to be trigonal planar, trigonal pyramidal or t-shaped, otherwise it assigned via hacks.
        elif (n_atoms == 4) and (central_element == "Cl") and (central_element == "Xe"):

            return "t_shaped"
        
        elif (n_atoms == 4) and (nuclear_charge >= 8):

            return "trigonal_pyramidal"
        
        elif (n_atoms == 4) and (nuclear_charge < 8):

            return "trigonal_planar"
        
        ################## n_adj=5 ##################
        
        elif (n_atoms == 5)  and (central_element == "C") and all_elements.count("C") == 1 and all_elements.count("O") == 2 and all_elements.count("H") == 1 and (adjacent_metal == True):

            return "formate"
        
        elif (n_atoms == 5)  and (central_element == "C") and all_elements.count("C") >= 3 and (adjacent_metal == True) and (is_in_cycle == True):

            return "ferrocene"
        
        elif (n_atoms == 5)  and (nuclear_charge < 28 or central_element == "Ge" or central_element == "Sn" or central_element == "Pb" or central_element == "As"):

            return "tetrahedral"

        elif (n_atoms == 5):

            return self._get_best_fit(soap, ["tetrahedral", "square_planar", "seesaw"])
        
        ################## n_adj=6 ##################

        elif (n_atoms == 6):

            return self._get_best_fit(soap, ["trigonal_bipyramidal", "square_pyramidal"])
        
        ################## n_adj>6 ##################

        elif (n_atoms == 7):

            return "octahedral"

        elif (n_atoms == 8):

            return "pentagonal_bipyramidal_formate"
        
        elif (n_atoms == 9):

            return "square_antiprismatic"
        
        else:
            
            return "unknown"

    def _get_best_fit(self, soap, possible_classes):

        # Calculate deviations from ideal local geometries
        deviations = {key: np.linalg.norm((self.ideal_soaps[key]-soap)) for key in possible_classes}

        # Return geometry class with smallest deviation
        chosen_class =  min(deviations, key=deviations.get)
        
        return chosen_class

    def _get_soap(self, coordinates, n_atoms, core_index):
        
        elements = ["H"]*n_atoms

        elements[core_index] = "C"

        soap = SOAP(
        species=["C", "H"],
        periodic=False,
        r_cut=3.0,
        n_max=2,
        l_max=1,
        )

        ase_atoms = Atoms("".join(elements), positions=coordinates)

        return soap.create(ase_atoms, centers=[core_index])


    def constrain_nodes_angles(self, is_adjacent_visible=False, include_all=False):

        with open(os.path.join(os.path.dirname(__file__), "..", "..", "data", 'ideal_geometries.json')) as json_file:

            self.ideal_geometries = json.load(json_file)

        if include_all:

            visible_nodes = self.nodes()

        else:

            visible_nodes = self.return_visible_nodes()

        self.set_nodes_visible(visible_nodes=visible_nodes)

        vsepr_classes = self.return_node_data(attribute="vsepr_classes", include_all=include_all)

        angle_data = []

        if self.constraints.get("constraints") != None:

            extra_constraints = self.constraints["constraints"].get("angles")

        else:

            extra_constraints = {}

        for node, vsepr_class in zip(visible_nodes, vsepr_classes):

            node_angle = self._constrain_node_angle(node, visible_nodes, vsepr_class, is_adjacent_visible=is_adjacent_visible, include_all=include_all)
   
            if extra_constraints and extra_constraints.get(node) != None:

                for extra_constraint in extra_constraints[node]:

                    node_angle += [extra_constraint]

                filtered_angles = {}

                for data_dict in node_angle:
                    nodes_key = tuple(data_dict["nodes"])
                    filtered_angles[nodes_key] = data_dict

                # Extract the filtered list of dictionaries containing the latest occurrences
                node_angle = list(filtered_angles.values())
        
            angle_data.append(node_angle)

        # Get new vsepr classes as they might have changed _constrain_node_angle

        self.set_nodes_visible(visible_nodes=visible_nodes)
        
        vsepr_classes = self.return_node_data(attribute="vsepr_classes", include_all=include_all)

        dihedral_data = [self._constrain_node_dihedral(node, visible_nodes, vsepr_class, is_adjacent_visible=is_adjacent_visible, include_all=include_all) for node, vsepr_class in zip(visible_nodes, vsepr_classes)]
  
        self.set_nodes_visible(visible_nodes=visible_nodes)

        self.add_node_data(attribute="angle_constraints", new_data=angle_data)

        self.add_node_data(attribute="dihedral_constraints", new_data=dihedral_data)

        del self.ideal_geometries


    def _constrain_node_angle(self, node, nodes, vsepr_class, is_adjacent_visible=False, include_all=False):

        if vsepr_class == "unknown" or vsepr_class == "dot" or vsepr_class == "linear_cn1":

            return []
        

        if vsepr_class in ["bent", "trigonal_planar", "trigonal_pyramidal", "tetrahedral"]:
        
            return self._constrain_simple_node_angles(node, nodes, vsepr_class,  is_adjacent_visible, include_all)
        
        self.set_nodes_visible(visible_nodes=nodes)

        adjacent_atoms = self.return_adjacent_nodes(node, is_adjacent_visible=is_adjacent_visible, include_all=include_all)

        self.set_nodes_visible(visible_nodes=[node]+adjacent_atoms)

        coordinates = self.return_node_data(attribute="coordinates")

        n_atoms = self.return_n_atoms()

        visible_nodes = self.return_visible_nodes()

        core_index = visible_nodes.index(node)

        if vsepr_class not in ["bent", "linear_cn2"]:

            best_order = self._get_best_order_for_match(n_atoms, coordinates, core_index, vsepr_class)
                   
        else:

            best_order = [core_index] + [i for i in range(0, n_atoms) if i != core_index]

        ideal_angles = self.ideal_geometries[vsepr_class]["angles"]

        final_angles = []

        for ideal_angle in ideal_angles:

            angle_nodes = [ideal_angle["nodes"][0], ideal_angle["nodes"][2]]

            angle_nodes_global = [visible_nodes[best_order[node]] for node in angle_nodes]

            final_angles.append({"nodes": angle_nodes_global, "angle": ideal_angle["angle"]*np.pi/180.})

        return final_angles
    
    def _get_best_order_for_match(self, n_atoms, coordinates, core_index, vsepr_class):

        coordinates = np.array([Vector(coordinates[[core_index, i]]).unit_vector() for i in range(0, n_atoms)])

        ideal_file_path = os.path.join(os.path.dirname(__file__), "..", "..", "data", "ideal_geometries", f'{vsepr_class}.xyz')

        _, ideal_coordinates, _ = FileReader(ideal_file_path).read_file()

        order = [i for i in range(0, n_atoms) if i != core_index]

        shuffled = list(permutations(order))

        rmsds = []

        for permutation in shuffled:

            coordinates_permuted = [coordinates[core_index]] + [coordinates[i] for i in permutation]

            _, rmsd = Rotation.align_vectors(coordinates_permuted, ideal_coordinates)

            rmsds.append(rmsd)

        best_order_index = np.argmin(rmsds)

        return [core_index] + list(shuffled[best_order_index])
    
    def _constrain_simple_node_angles(self, node, nodes, vsepr_class, is_adjacent_visible, include_all):

        self.set_nodes_visible(visible_nodes=node)

        cycle_size = self.return_node_data("smallest_cycle_size")[0]

        cycle_ids_of_node = self.return_node_data(attribute="cycle_id")[0]

        central_node_in_cycle = cycle_ids_of_node is not None

        if central_node_in_cycle:

            n_cycles = len(cycle_ids_of_node)

        else:

            n_cycles = 0

        self.set_nodes_visible(visible_nodes=nodes)

        adjacent_atoms = self.return_adjacent_nodes(node, is_adjacent_visible=is_adjacent_visible, include_all=include_all)

        self.set_nodes_visible(visible_nodes=[node]+adjacent_atoms)

        visible_nodes = self.return_visible_nodes()

        cycle_ids_of_nodes = self.return_node_data(attribute="cycle_id")

        core_index = visible_nodes.index(node)

        if central_node_in_cycle:

            central_cycle_id = cycle_ids_of_node[0]

            local_non_cycle_nodes = [i for i, indices in enumerate(cycle_ids_of_nodes) if i != core_index and (indices is None or central_cycle_id not in indices)]

        else:

            local_non_cycle_nodes = []

        if central_node_in_cycle and n_cycles == 1 and cycle_size == 3 and vsepr_class in ["tetrahedral", "bent"]:
        
            final_angles = []

            if vsepr_class == "tetrahedral":
                
                global_non_cycle_nodes = [visible_nodes[i] for i in local_non_cycle_nodes]
            
                global_cycle_nodes = [i for i in visible_nodes if i not in global_non_cycle_nodes and i != node]

                for global_cycle_node in global_cycle_nodes:

                    for global_non_cycle_node in global_non_cycle_nodes:

                        final_angles.append({"nodes": [global_non_cycle_node, global_cycle_node], "angle": 120.*np.pi/180.})

                final_angles.append({"nodes": global_cycle_nodes, "angle": 60.0*np.pi/180.})

                final_angles.append({"nodes": global_non_cycle_nodes, "angle": 109.0*np.pi/180.})

                return final_angles
            
            else:

                return [{"nodes": adjacent_atoms, "angle": 60.*np.pi/180.}]
            
        elif central_node_in_cycle and n_cycles == 1 and vsepr_class == "trigonal_planar" and len(local_non_cycle_nodes) == 1:
                
            final_angles = []

            if cycle_size == 3:

                major_angle = 150.0

                minor_angle = 60.0

            elif cycle_size == 4:

                major_angle = 135.0

                minor_angle = 90.0

            elif cycle_size == 5:

                major_angle = 126.0

                minor_angle = 108.0

            else:

                major_angle = 120.0

                minor_angle = 120.0

            global_non_cycle_node = [visible_nodes[i] for i in local_non_cycle_nodes][0]
        
            global_cycle_nodes = [i for i in visible_nodes if i != global_non_cycle_node and i != node]

            for cycle_node in global_cycle_nodes:

                final_angles.append({"nodes": [global_non_cycle_node, cycle_node], "angle": major_angle*np.pi/180.})

            final_angles.append({"nodes": global_cycle_nodes, "angle": minor_angle*np.pi/180.})

            return final_angles
        
        elif vsepr_class == "trigonal_pyramidal" and central_node_in_cycle == False:

                local_priorities = self.return_node_data(attribute="priorities")

                adjacent_priorities = [local_priorities[visible_nodes.index(i)] for i in adjacent_atoms]

                if len(set(adjacent_priorities)) != 3:

                    self.nodes[node]["vsepr_classes"] = "trigonal_planar"

                    vsepr_class == "trigonal_planar"

        coordinates = self.return_node_data(attribute="coordinates")

        n_atoms = self.return_n_atoms()

        local_angle_nodes = [[i, core_index, j] for i in range(n_atoms) if i != core_index for j in range(n_atoms) if j != core_index  and j > i]

        global_angle_nodes = [[visible_nodes[angle_nodes[0]],visible_nodes[angle_nodes[2]] ]for angle_nodes in local_angle_nodes]

        real_angles = [Vector(coordinates[angle]).calculate_angle() for angle in local_angle_nodes] 

        ideal_angles = list(set([angle["angle"] for angle in self.ideal_geometries[vsepr_class]["angles"]]))

        final_angles = []

        for angle_nodes, real_angle in zip(global_angle_nodes, real_angles):

            deviations = [np.abs(real_angle-ideal_angle) for ideal_angle in ideal_angles]

            argmin_dev = np.argmin(deviations)

            final_angles.append({"nodes": angle_nodes, "angle": ideal_angles[argmin_dev]*np.pi/180.})

        return final_angles

    def _constrain_node_dihedral(self, node, nodes, vsepr_class, is_adjacent_visible=False, include_all=False):

        if vsepr_class not in ["tetrahedral", "trigonal_pyramidal", "trigonal_planar", "square_planar", "formate"]:

            return []
        
        self.set_nodes_visible(visible_nodes=nodes)

        adjacent_nodes = self.return_adjacent_nodes(node, is_adjacent_visible=is_adjacent_visible, include_all=include_all)

        self.set_nodes_visible(visible_nodes=[node]+adjacent_nodes)

        visible_nodes = self.return_visible_nodes()

        core_index = visible_nodes.index(node)

        if vsepr_class == "tetrahedral":

            return self._constrain_squezzed_tetrahedral_dihedral(core_index, adjacent_nodes, visible_nodes)

        elif vsepr_class == "trigonal_pyramidal":

            return self._constrain_pyramidal_dihedral(core_index, visible_nodes)  
        
        else:
  
            return self._constrain_planar_dihedral(core_index, visible_nodes)


    def _constrain_squezzed_tetrahedral_dihedral(self, core_index: int, adjacent_nodes: np.ndarray, visible_nodes: np.ndarray):
        """
        Constrain the dihedral angles of a node with tetrahedral geometry if node is part of two cycles and the tetrahedral geometry is far from ideal. 
        Prevents flip of chirality.

        Returns: 
            list: list of dictionaries containing the nodes and the ideal dihedral angle
        """
        priorities = self.return_node_data(attribute="priorities")

        cycle_ids = self.return_node_data(attribute="cycle_id")

        adjacent_priorities = [priorities[visible_nodes.index(i)] for i in adjacent_nodes]

        # If tetrahedral node is achiral return no dihedral constraints
        if len(set(adjacent_priorities)) != 4:


            return []

        core_cycle_id = cycle_ids[core_index]

        # If tetrahedral node is not part of two cycles return no dihedral constraints as otherwise it is constrained enough
        if core_cycle_id is None or len(core_cycle_id) < 2:

            return []
        

        adjacent_cycle_ids = [cycle_ids[visible_nodes.index(i)] for i in adjacent_nodes]

        sum_adjacent_not_in_cycle = sum(1 for cycle_id in adjacent_cycle_ids if cycle_id is None)

        # If no adjacent is no part of a cycle return no dihedral constraints as otherwise it is constrained enough.
        if sum_adjacent_not_in_cycle != 1:

            return []
        
        cycles = [self.cycle_nodes[idx] for idx in core_cycle_id]

        len_cycles = [len(cycle) for cycle in cycles]

        max_len = max(len_cycles)

        non_cycle_node = visible_nodes.index(adjacent_nodes[[i for i, cycle_id in enumerate(adjacent_cycle_ids) if cycle_id is None][0]])

        #cycles_of_interest = [cycle for cycle in cycles if len(cycle) == max_len]

        cycles_of_interest = cycles.copy()

        final_dihedrals = []

        # For each cycle of interest, the dihedral angle is constrained to be close to 120 or -120 degrees
        for cycle_of_interest in cycles_of_interest:

            dihedrals_for_cycle = self.constrain_tetrahedral_dihedral_for_one_cycle(core_index, non_cycle_node, priorities, cycle_of_interest, adjacent_nodes, visible_nodes)

            final_dihedrals += dihedrals_for_cycle

        return final_dihedrals
    

    def constrain_tetrahedral_dihedral_for_one_cycle(self, core_index: int, non_cycle_node: int, priorities: np.ndarray, cycle_of_interest: np.ndarray, adjacent_nodes: np.ndarray, visible_nodes: np.ndarray):
        """
        Get dihedral angles for one cycle for a node with tetrahedral geometry if node is part of two cycles and the tetrahedral geometry is far from ideal. 

        Args:
            core_index (int): index of the central node in visible_nodes
            non_cycle_node (int): index of the non-cycle node in visible_nodes
            priorities (np.ndarray): priorities of all visible nodes
            cycle_of_interest (np.ndarray): cycle of interest
            adjacent_nodes (np.ndarray): adjacent nodes of the central node
            visible_nodes (np.ndarray): visible nodes of the molecule
        """

        adjacent_node_in_larger_cycle = [visible_nodes.index(i) for i in adjacent_nodes if i in cycle_of_interest]

        priorities_in_larger_cycle = [priorities[i] for i in adjacent_node_in_larger_cycle]

        max_priority = np.argmax(priorities_in_larger_cycle)

        min_priority = np.argmin(priorities_in_larger_cycle)

        node_with_max_priority = adjacent_node_in_larger_cycle[max_priority]

        node_with_min_priority = adjacent_node_in_larger_cycle[min_priority]

        dihedrals = [[node_with_min_priority, core_index, node_with_max_priority, non_cycle_node], [node_with_max_priority, core_index, node_with_min_priority, non_cycle_node]]

        coordinates = self.return_node_data(attribute="coordinates")

        final_dihedrals = []

        for dihedral in dihedrals:

            real_dihedral = Vector(coordinates[dihedral]).calculate_dihedral()

            ideal_dihedrals = [120.0, -120.0]

            deviations = [np.abs(real_dihedral-ideal_dihedral) for ideal_dihedral in ideal_dihedrals]

            argmin_dev = np.argmin(deviations)
            
            final_dihedrals.append({"nodes": [visible_nodes[dihedral[0]], visible_nodes[dihedral[2]], visible_nodes[dihedral[3]]], "angle": ideal_dihedrals[argmin_dev]*np.pi/180.})

        return final_dihedrals
    
    def _constrain_pyramidal_dihedral(self, core_index, visible_nodes):

        ideal_dihedrals = [120.0, -120.0]

        n_atoms = self.return_n_atoms()

        coordinates = self.return_node_data(attribute="coordinates")

        dihedral_nodes = self._get_dihedral_nodes(coordinates, n_atoms, core_index)

        real_dihedrals = [Vector(coordinates[dihedral]).calculate_dihedral() for dihedral in dihedral_nodes] 

        assigned_dihedrals = []

        for real_dihedral in real_dihedrals:

            deviations = []

            for ideal_dihedral in ideal_dihedrals:

                deviation = np.abs(ideal_dihedral-real_dihedral)

                if deviation < 90.0:

                    deviations.append(deviation)

                else:

                    deviations.append(np.inf)

            smallest_deviation = np.argmin(deviations)

            if deviations[smallest_deviation] != np.inf:

                assigned_dihedrals.append(ideal_dihedrals[smallest_deviation])

            else:
                    
                assigned_dihedrals.append(np.inf)

        return [{"nodes": [visible_nodes[dihedral[0]], visible_nodes[dihedral[2]], visible_nodes[dihedral[3]]], "angle": assigned_angle*np.pi/180.} for dihedral, assigned_angle in zip(dihedral_nodes, assigned_dihedrals) if assigned_angle != np.inf]
    

    def _constrain_planar_dihedral(self, core_index, visible_nodes):

        ideal_dihedrals = [180.0, -180.0]

        n_atoms = self.return_n_atoms()

        coordinates = self.return_node_data(attribute="coordinates")

        dihedral_nodes = self._get_dihedral_nodes(coordinates, n_atoms, core_index)

        real_dihedrals = [Vector(coordinates[dihedral]).calculate_dihedral() for dihedral in dihedral_nodes] 

        assigned_dihedrals = []

        for real_dihedral in real_dihedrals:

            deviations = []

            for ideal_dihedral in ideal_dihedrals:

                deviation = np.abs(ideal_dihedral-real_dihedral)

                #Dummy value, alway flatten, no matter the deviation
                if deviation < 180.0:

                    deviations.append(deviation)

                else:

                    deviations.append(np.inf)

            smallest_deviation = np.argmin(deviations)

            if deviations[smallest_deviation] != np.inf:

                assigned_dihedrals.append(ideal_dihedrals[smallest_deviation])

            else:
                    
                assigned_dihedrals.append(np.inf)

        return [{"nodes": [visible_nodes[dihedral[0]], visible_nodes[dihedral[2]], visible_nodes[dihedral[3]]], "angle": assigned_angle*np.pi/180.} for dihedral, assigned_angle in zip(dihedral_nodes, assigned_dihedrals) if assigned_angle != np.inf]
    

    def _get_dihedral_nodes(self, coordinates, n_atoms, core_index):
        
        #dihedral_nodes =  [[i, core_index, j, k] for i in range(n_atoms) for j in range(n_atoms) for k in range(n_atoms) if len(set([i, core_index, j, k])) == len([i, core_index, j, k]) and k > i and i != core_index and j != core_index and k != core_index]

        dihedral_nodes =  [[i, core_index, j, k] for i in range(n_atoms) for j in range(n_atoms) for k in range(n_atoms) if len(set([i, core_index, j, k])) == len([i, core_index, j, k])]

        final_dihedral_nodes = []

        for dihedral in dihedral_nodes:

            angle1 = Vector(coordinates[[dihedral[0], dihedral[1], dihedral[2]]]).calculate_angle()

            angle2 = Vector(coordinates[[dihedral[2], dihedral[1], dihedral[3]]]).calculate_angle()

            if (angle1 < 150.0 or angle1 > 210.0) and (angle2 < 150.0 or angle2 > 210.0):

                final_dihedral_nodes.append(dihedral)

        return final_dihedral_nodes
        

