"""
Module for handling molecules. Uses the pymatgen.core.structure.Molecule
class as a base. Has a function for reorienting molecules
(reoriented_molecule), and for calculating valid orientations within a Wyckoff
position based on symmetry (orientation_in_wyckoff_position).
The orientation class can be used to identify
degrees of freedom for molecules in Wyckoff positions with certain symmetry
constraints.

"""
# Imports
import numpy as np
from copy import deepcopy
from scipy.spatial.transform import Rotation
import networkx as nx

# ------------------------------
# External Libraries
from pymatgen.core.structure import Molecule
from pymatgen.symmetry.analyzer import PointGroupAnalyzer, generate_full_symmops
from pymatgen.core.bonds import CovalentBond

# PyXtal imports
from pyxtal.msg import printx
from pyxtal.tolerance import Tol_matrix
from pyxtal.database.element import Element
from pyxtal.operations import SymmOp, OperationAnalyzer, rotate_vector, angle
from pyxtal.database.collection import Collection

# Define functions
# ------------------------------
molecule_collection = Collection("molecules")


class pyxtal_molecule:
    """
    Extended molecule class based on pymatgen.core.structure.Molecule
    The added features include:
    0, parse the input
    1, estimate volume/tolerance/radii
    2, find and store symmetry (todo)
    3, get the principle axis (todo)
    4, re-align the molecule (todo)

    Args:
        mol: a string to reprent the molecule
        tm: tolerance matrix
   
    """

    def __init__(self, mol, tm=Tol_matrix(prototype="molecular")):
        mo = None
        if type(mol) == str:
            # Parse molecules: either file or molecule name
            tmp = mol.split(".")
            if len(tmp) > 1:
                # Load the molecule from the given file
                if tmp[-1] in ["xyz", "gjf", "g03", "json"]:
                    import os
                    if os.path.exists(mol):
                        mo = Molecule.from_file(mol)
                    else:
                        raise NameError("{:s} is not a valid path".format(mol))
                else:
                    raise NameError("{:s} is not a supported format".format(tmp[-1]))
            else:
                # print('\nLoad the molecule {:s} from collections'.format(mol))
                mo = molecule_collection[mol]
        elif hasattr(mol, "sites"):  # pymatgen molecule
            mo = mol

        if mo is None:
            msg = "Could not create molecules from given input: {:s}".format(mol)
            raise NameError(msg)

        self.props = mo.site_properties

        if len(mo) > 1:
            pga = PointGroupAnalyzer(mo)
            mo = pga.symmetrize_molecule()["sym_mol"] 
            mo = self.add_site_props(mo)

        self.mol = mo
        self.tm = tm
        self.get_box()
        self.volume = self.box.volume
        self.get_radius()
        self.get_symbols()
        self.get_tols_matrix()

    def swap_axis(self, ax):
        """
        swap the molecular axis
        """
        coords = self.mol.cart_coords[:, ax]
        mo = Molecule(self.symbols, coords)
        mo = self.add_site_props(mo)

        return pyxtal_molecule(mo, self.tm)


    def add_site_props(self, mo):
        if len(self.props) > 0:
            for key in self.props.keys():
                mo.add_site_property(key, self.props[key])
        return mo

    def get_box(self):
        """
        Given a molecule, find a minimum orthorhombic box containing it.
        Size is calculated using min and max x, y, and z values, 
        plus the padding defined by the vdw radius
        For best results, call oriented_molecule first.
        
        Args:
            mol: a pymatgen Molecule object. Should be oriented along its principle axes.
    
        Returns:
            a Box object
        """
        mol, _ = reoriented_molecule(self.mol)
        minx, miny, minz, maxx, maxy, maxz = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
        for p in mol:
            x, y, z = p.coords
            r = Element(p.species_string).vdw_radius
            if x - r < minx:
                minx = x - r
            if y - r < miny:
                miny = y - r
            if z - r < minz:
                minz = z - r
            if x + r > maxx:
                maxx = x + r
            if y + r > maxy:
                maxy = y + r
            if z + r > maxz:
                maxz = z + r
        self.box = Box(minx, maxx, miny, maxy, minz, maxz)

    def get_radius(self):
        r_max = 0
        for coord, number in zip(self.mol.cart_coords, self.mol.atomic_numbers):
            radius = (
                np.sqrt(np.sum(coord * coord)) + self.tm.get_tol(number, number) * 0.5
            )
            if radius > r_max:
                r_max = radius
        self.radius = r_max

    def get_symbols(self):
        self.symbols = [specie.name for specie in self.mol.species]

    def get_tols_matrix(self):
        """
        Returns: a 2D matrix which is used internally for distance checking.
        """
        numbers = self.mol.atomic_numbers
        tols = np.zeros((len(numbers), len(numbers)))
        for i1, number1 in enumerate(numbers):
            for i2, number2 in enumerate(numbers):
                tols[i1][i2] = self.tm.get_tol(number1, number2)
        if len(self.mol)==1:
            tols *= 0.8 # if only one atom, reduce the tolerance
        self.tols_matrix = tols

    def show(self):
        from pyxtal.viz import display_molecules
        return display_molecules([self.mol])




class Box:
    """
    Class for storing the binding box for a molecule. Box is oriented along the x, y, and
    z axes.

    Args:
        minx: the minimum x value
        maxx: the maximum x value
        miny: the minimum y value
        maxy: the maximum y value
        minz: the minimum z value
        maxz: the maximum z value
    """

    def __init__(self, minx, maxx, miny, maxy, minz, maxz):
        self.minx = float(minx)
        self.maxx = float(maxx)
        self.miny = float(miny)
        self.maxy = float(maxy)
        self.minz = float(minz)
        self.maxz = float(maxz)

        self.width = float(abs(maxx - minx))
        self.length = float(abs(maxy - miny))
        self.height = float(abs(maxz - minz))

        self.minl = min(self.width, self.length, self.height)
        self.maxl = max(self.width, self.length, self.height)
        for x in (self.width, self.length, self.height):
            if x <= self.maxl and x >= self.minl:
                self.midl = x

        self.volume = float(self.width * self.length * self.height)


class Orientation:
    """
    Stores orientations for molecules based on vector constraints.
    Can be stored to regenerate orientations consistent with a given constraint
    vector, without re-calling orientation_in_wyckoff_position. Allows for
    generating orientations which differ only in their rotation about a given
    axis.

    Args:
        matrix: a 3x3 rotation matrix describing the orientation (and/or
            inversion) to store
        degrees: the number of degrees of freedom...
            0 - The orientation refers to a single rotation matrix
            1 - The orientation can be rotated about a single axis
            2 - The orientation can be any pure rotation matrix

        axis:
            an optional axis about which the orientation can rotate. Only used
            if degrees is equal to 1
    """

    def __init__(self, matrix, degrees=2, axis=None):
        self.matrix = np.array(matrix)
        self.r = Rotation.from_matrix(matrix)  # scipy transform.Rotation class
        self.degrees = degrees  # The number of degrees of freedom.
        if degrees == 1:
            if axis is None:
                raise ValueError("axis is required for orientation")
            else:
                axis /= np.linalg.norm(axis)
        self.axis = axis
        self.angle = None

    def __str__(self):
        s = "-------PyXtal.molecule.Orientation class----\n"
        s += "degree of freedom: {:d}\n".format(self.degrees)
        s += "Rotation matrix:\n"
        s += "{:6.3f} {:6.3f} {:6.3f}\n".format(*self.matrix[:,0])
        s += "{:6.3f} {:6.3f} {:6.3f}\n".format(*self.matrix[:,1])
        s += "{:6.3f} {:6.3f} {:6.3f}\n".format(*self.matrix[:,2])
        if self.axis is not None:
            s += "Rotation axis\n"
            s += "{:6.2f} {:6.2f} {:6.3f}\n".format(*self.axis)
        return s

    def __repr__(self):
        return str(self)


    def copy(self):
        return deepcopy(self)

    def change_orientation(self, angle="random"):
        """
        Allows for specification of an angle (possibly random) to
        rotate about the constraint axis.

        Args:
            angle: an angle to rotate about the constraint axis. 
            If "random", chooses a random rotation angle. 
            If self.degrees==2, chooses a random rotation matrix. 
            If self.degrees==1, only apply on angle
            If self.degrees==0, no change

        """
        if self.degrees >= 1:
            # choose the axis
            if self.axis is None:
                axis = np.random.RandomState().rand(3) - 0.5
                self.axis = axis / np.linalg.norm(axis)
 
            # parse the angle
            if angle == "random":
                angle = np.random.RandomState().rand() * np.pi * 2
            self.angle = angle
    
            # update the matrix
            r1 = Rotation.from_rotvec(self.angle * self.axis)
            self.r = r1 * self.r
            #self.r *= r1 
            self.matrix = self.r.as_matrix()

    def rotate_by_matrix(self, matrix, ignore_constraint=True):
        """
        rotate

        Args:
            matrix: 3*3 rotation matrix

        """
        if not ignore_constraint:
            if self.degrees == 0:
                raise ValueError("cannot rotate")
            elif self.degrees == 1:
                axis = self.axis
                vec = Rotation.from_matrix(matrix).as_rotvec()
                if angle(vec, self.axis) > 1e-2 and angle(vec, -self.axis) > 1e-2:
                    raise ValueError("must rotate along the given axis")
        else:
            axis = None
            
        matrix = matrix.dot(self.matrix)
        return Orientation(matrix, self.degrees, axis)

    def get_matrix(self, angle="random"):
        """
        Generate a 3x3 rotation matrix consistent with the orientation's
        constraints. Allows for specification of an angle (possibly random) to
        rotate about the constraint axis.

        Args:
            angle: an angle to rotate about the constraint axis. If "random",
                chooses a random rotation angle. If self.degrees==2, chooses a
                random 3d rotation matrix to multiply by. If the original matrix
                is wanted, set angle=0, or call self.matrix

        Returns:
            a 3x3 rotation (and/or inversion) matrix (numpy array)
        """
        if self.degrees == 2:
            if angle == "random":
                axis = np.random.sample(3)
                axis = axis / np.linalg.norm(axis)
                angle = np.random.random() * np.pi * 2
            else:
                axis = self.axis
            return Rotation.from_rotvec(angle * axis).as_matrix()

        elif self.degrees == 1:
            if angle == "random":
                angle = np.random.random() * np.pi * 2
            else:
                angle = self.angle
            return Rotation.from_rotvec(angle * self.axis).as_matrix()

        elif self.degrees == 0:
            return self.matrix

    def get_op(self, angle=None):
        """
        Generate a SymmOp object consistent with the orientation's
        constraints. Allows for specification of an angle (possibly random) to
        rotate about the constraint axis.

        Args:
            angle: an angle to rotate about the constraint axis. If "random",
                chooses a random rotation angle. If self.degrees==2, chooses a
                random 3d rotation matrix to multiply by. If the original matrix
                is wanted, set angle=0, or call self.matrix

        Returns:
            pymatgen.core.structure. SymmOp object
        """
        #if angle is not None:
        #    self.change_orientation(angle)
        return SymmOp.from_rotation_and_translation(self.matrix, [0, 0, 0])

    @classmethod
    def from_constraint(self, v1, c1):
        """
        Geneate an orientation object given a constraint axis c1, and a
        corresponding vector v1. v1 will be rotated onto c1, and the resulting
        orientation will have a rotational degree of freedom about c1.

        Args:
            v1: a 1x3 vector in the original reference frame
            c1: a corresponding axis which v1 must be mapped to

        Returns:
            an orientation object consistent with the supplied constraint
        """
        # c1 is the constraint vector; v1 will be rotated onto it
        m = rotate_vector(v1, c1)
        return Orientation(m, degrees=1, axis=c1)

    @classmethod
    def from_constraints(self, v1, c1, v2, c2):
        """
        Geneate an orientation object given two constraint vectors

        Args:
            v1: a 1x3 vector in the original reference frame
            c1: a corresponding axis which v1 must be mapped to
            v1: a second 1x3 vector in the original reference frame
            c1: a corresponding axis which v2 must be mapped to

        Returns:
            an orientation object consistent with the supplied constraints
        """
        T = rotate_vector(v1, c1)
        phi = angle(c1, c2)
        phi2 = angle(c1, (np.dot(T, v2)))
        if not np.isclose(phi, phi2, rtol=0.01):
            printx("Error: constraints and vectors do not match.", priority=1)
            return
        r = np.sin(phi)
        c = np.linalg.norm(np.dot(T, v2) - c2)
        theta = np.arccos(1 - (c ** 2) / (2 * (r ** 2)))
        Rot = R.from_rotvec(theta * c1)
        T2 = np.dot(R, T)
        a = angle(np.dot(T2, v2), c2)
        if not np.isclose(a, 0, rtol=0.01):
            T2 = np.dot(np.linalg.inv(R), T)
        a = angle(np.dot(T2, v2), c2)
        if not np.isclose(a, 0, rtol=0.01):
            printx("Error: Generated incorrect rotation: " + str(theta), priority=1)
        return Orientation(T2, degrees=0)

    def random_orientation(self):
        """
        Applies random rotation (if possible) and returns a new orientation with
        the new base matrix.

        Returns:
            a new orientation object with a different base rotation matrix
        """

        self.change_orientation()
        return self

    def get_Euler_angles(self):
        return self.r.as_euler('zxy', degrees=True)


def get_inertia_tensor(mol):
    """
    Calculate the symmetric inertia tensor for a Molecule object. Used to find
    the principal axes of symmetry.

    Args:
        mol: a Molecule object

    Returns:
        a 3x3 numpy array representing a moment of inertia tensor
    """
    mo = mol.get_centered_molecule()
    # Initialize elements of the inertia tensor
    I11 = I22 = I33 = I12 = I13 = I23 = 0.0
    for i in range(len(mo)):
        x, y, z = mo.cart_coords[i]
        m = 1 #mo[i].specie.number
        I11 += m * (y ** 2 + z ** 2)
        I22 += m * (x ** 2 + z ** 2)
        I33 += m * (x ** 2 + y ** 2)
        I12 += -m * x * y
        I13 += -m * x * z
        I23 += -m * y * z
    return np.array([[I11, I12, I13], [I12, I22, I23], [I13, I23, I33]])


def reoriented_molecule(mol, nested=False):
    """
    Reorient a molecule so that its principal axes are aligned with the
    identity matrix.

    Args:
        mol: a Molecule object
        nested: internal variable to keep track of how many times the function
            has been called recursively

    Returns:
        new_mol: a reoriented copy of the original molecule. 
        P: the 3x3 rotation matrix used to obtain it.
    """

    def reorient(mol):
        new_mol = mol.get_centered_molecule()
        A = get_inertia_tensor(new_mol)
        # Store the eigenvectors of the inertia tensor
        P = np.transpose(np.linalg.eigh(A)[1])
        if np.linalg.det(P) < 0:
            P[0] *= -1
        # reorient the molecule
        P = SymmOp.from_rotation_and_translation(P, [0, 0, 0])
        new_mol.apply_operation(P)
        # Our molecule should never be inverted during reorientation.
        if np.linalg.det(P.rotation_matrix) < 0:
            printx(
                "Error: inverted reorientation applied.\n"
                + "(Within reoriented_molecule)",
                priority=0,
            )
        return new_mol, P

    # If needed, recursively apply reorientation (due to numerical errors)
    iterations = 1
    max_iterations = 100
    new_mol, P = reorient(mol)
    while iterations < max_iterations:
        is_okay = True
        for i in range(3):
            for j in range(3):
                x = np.linalg.eigh(get_inertia_tensor(new_mol))[1][i][j]
                okay = True
                if i == j:
                    # Check that diagonal elements are 0 or 1
                    if (not np.isclose(x, 0)) and (not np.isclose(x, 1)):
                        okay = False
                else:
                    # Check that off-diagonal elements are 0
                    if not np.isclose(x, 0):
                        okay = False
        if okay is False:
            # If matrix is not diagonal with 1's and/or 0's, reorient
            new_mol, Q = reorient(new_mol)
            P = Q * P
            iterations += 1
        elif okay is True:
            break
    if iterations == max_iterations:
        printx("Error: Could not reorient molecule after ", priority=0)
        return False
    return new_mol, P


def get_symmetry(mol, already_oriented=False):
    """
    Return a molecule's point symmetry.
    Note: for linear molecules, infinitessimal rotations are treated as 6-fold
    rotations, which works for 3d and 2d point groups.

    Args:
        mol: a Molecule object
        already_oriented: whether or not the principle axes of mol are already
            reoriented. Can save time if True, but is not required.

    Returns:
        a list of SymmOp objects which leave the molecule unchanged when applied
    """
    # For single atoms, we cannot represent the point group using a list of operations
    if len(mol) == 1:
        return []
    pga = PointGroupAnalyzer(mol)
    # Handle linear molecules
    if "*" in pga.sch_symbol:
        if already_oriented == False:
            # Reorient the molecule
            oriented_mol, P = reoriented_molecule(mol)
            pga = PointGroupAnalyzer(oriented_mol)
        pg = pga.get_pointgroup()
        symm_m = []
        for op in pg:
            symm_m.append(op)
        # Add 12-fold  and reflections in place of ininitesimal rotation
        for axis in [[1, 0, 0], [0, 1, 0], [0, 0, 1]]:
            # op = SymmOp.from_rotation_and_translation(aa2matrix(axis, np.pi/6), [0,0,0])
            m1 = Rotation.from_rotvec(np.pi / 6 * axis).as_matrix()
            op = SymmOp.from_rotation_and_translation(m1, [0, 0, 0])
            if pga.is_valid_op(op):
                symm_m.append(op)
                # Any molecule with infinitesimal symmetry is linear;
                # Thus, it possess mirror symmetry for any axis perpendicular
                # To the rotational axis. pymatgen does not add this symmetry
                # for all linear molecules - for example, hydrogen
                if axis == [1, 0, 0]:
                    symm_m.append(SymmOp.from_xyz_string("x,-y,z"))
                    symm_m.append(SymmOp.from_xyz_string("x,y,-z"))
                    r = SymmOp.from_xyz_string("-x,y,-z")
                elif axis == [0, 1, 0]:
                    symm_m.append(SymmOp.from_xyz_string("-x,y,z"))
                    symm_m.append(SymmOp.from_xyz_string("x,y,-z"))
                    r = SymmOp.from_xyz_string("-x,-y,z")
                elif axis == [0, 0, 1]:
                    symm_m.append(SymmOp.from_xyz_string("-x,y,z"))
                    symm_m.append(SymmOp.from_xyz_string("x,-y,z"))
                    r = SymmOp.from_xyz_string("x,-y,-z")
                # Generate a full list of SymmOps for the molecule's pointgroup
                symm_m = generate_full_symmops(symm_m, 1e-3)
                break
        # Reorient the SymmOps into mol's original frame
        if not already_oriented:
            new = []
            for op in symm_m:
                new.append(P.inverse * op * P)
            return new
        elif already_oriented:
            return symm_m
    # Handle nonlinear molecules
    else:
        pg = pga.get_pointgroup()
        symm_m = []
        for op in pg:
            symm_m.append(op)
        return symm_m


def orientation_in_wyckoff_position(
    mol,
    wyckoff_position,
    randomize=True,
    exact_orientation=False,
    already_oriented=False,
    allow_inversion=True,
    rtol = 1e-2,
):
    """
    Tests if a molecule meets the symmetry requirements of a Wyckoff position,
    and returns the valid orientations.

    Args:
        mol: a Molecule object. Orientation is arbitrary
        wyckoff_position: a pyxtal.symmetry.Wyckoff_position object
        randomize: whether or not to apply a random rotation consistent with
            the symmetry requirements
        exact_orientation: whether to only check compatibility for the provided
            orientation of the molecule. Used within general case for checking.
            If True, this function only returns True or False
        already_oriented: whether or not to reorient the principle axes
            when calling get_symmetry. Setting to True can remove redundancy,
            but is not necessary
        allow_inversion: whether or not to allow chiral molecules to be
            inverted. Should only be True if the chemical and biological
            properties of the mirror image are known to be suitable for the
            desired application

    Returns:
        a list of operations.Orientation objects which can be applied to the
        molecule while allowing it to satisfy the symmetry requirements of the
        Wyckoff position. If no orientations are found, returns False.
    """
    # For single atoms, there are no constraints
    if len(mol) == 1:
        return [Orientation([[1, 0, 0], [0, 1, 0], [0, 0, 1]], degrees=2)]

    wyckoffs = wyckoff_position.ops
    w_symm = wyckoff_position.symmetry_m
    index = wyckoff_position.index

    # Obtain the Wyckoff symmetry
    symm_w = w_symm[0]
    pga = PointGroupAnalyzer(mol)

    # Check exact orientation
    if exact_orientation is True:
        mo = deepcopy(mol)
        valid = True
        for op in symm_w:
            if not pga.is_valid_op(op):
                valid = False
        if valid is True:
            return True
        elif valid is False:
            return False

    # Obtain molecular symmetry, exact_orientation==False
    symm_m = get_symmetry(mol, already_oriented=already_oriented)
    # Store OperationAnalyzer objects for each molecular SymmOp
    chiral = True
    opa_m = []
    for op_m in symm_m:
        opa = OperationAnalyzer(op_m)
        opa_m.append(opa)
        if opa.type == "rotoinversion":
            chiral = False
        elif opa.type == "inversion":
            chiral = False

    # If molecule is chiral and allow_inversion is False,
    # check if WP breaks symmetry
    if chiral is True:
        if allow_inversion is False:
            for op in wyckoffs:
                if np.linalg.det(op.rotation_matrix) < 0:
                    printx(
                        "Warning: cannot place chiral molecule in spagegroup", priority=2,
                    )
                    return False

    # Store OperationAnalyzer objects for each Wyckoff symmetry SymmOp
    opa_w = []
    for op_w in symm_w:
        opa_w.append(OperationAnalyzer(op_w))

    # Check for constraints from the Wyckoff symmetry...
    # If we find ANY two constraints (SymmOps with unique axes), the molecule's
    # point group MUST contain SymmOps which can be aligned to these particular
    # constraints. However, there may be multiple compatible orientations of the
    # molecule consistent with these constraints
    constraint1 = None
    constraint2 = None
    for i, op_w in enumerate(symm_w):
        if opa_w[i].axis is not None:
            constraint1 = opa_w[i]
            for j, op_w in enumerate(symm_w):
                if opa_w[j].axis is not None:
                    dot = np.dot(opa_w[i].axis, opa_w[j].axis)
                    if (not np.isclose(dot, 1, rtol=rtol)) and (
                        not np.isclose(dot, -1, rtol=rtol)
                    ):
                        constraint2 = opa_w[j]
                        break
            break
    # Indirectly store the angle between the constraint axes
    if constraint1 is not None and constraint2 is not None:
        dot_w = np.dot(constraint1.axis, constraint2.axis)
    # Generate 1st consistent molecular constraints
    constraints_m = []
    if constraint1 is not None:
        for i, opa1 in enumerate(opa_m):
            if opa1.is_conjugate(constraint1):
                constraints_m.append([opa1, []])
                # Generate 2nd constraint in opposite direction
                extra = deepcopy(opa1)
                extra.axis = [opa1.axis[0] * -1, opa1.axis[1] * -1, opa1.axis[2] * -1]
                constraints_m.append([extra, []])

    # Remove redundancy for the first constraints
    list_i = list(range(len(constraints_m)))
    list_j = list(range(len(constraints_m)))
    copy = deepcopy(constraints_m)
    for i, c1 in enumerate(copy):
        if i in list_i:
            for j, c2 in enumerate(copy):
                if i > j and j in list_j and j in list_i:
                    # Check if axes are colinear
                    if np.isclose(np.dot(c1[0].axis, c2[0].axis), 1, rtol=rtol):
                        list_i.remove(j)
                        list_j.remove(j)
                    # Check if axes are symmetrically equivalent
                    else:
                        cond1 = False
                        cond2 = False
                        for opa in opa_m:
                            if opa.type == "rotation":
                                op = opa.op
                                if np.isclose(
                                    np.dot(op.operate(c1[0].axis), c2[0].axis),
                                    1,
                                    rtol=5*rtol,
                                ):
                                    cond1 = True
                                    break
                        if cond1 is True:  # or cond2 is True:
                            list_i.remove(j)
                            list_j.remove(j)
    c_m = deepcopy(constraints_m)
    constraints_m = []
    for i in list_i:
        constraints_m.append(c_m[i])

    # Generate 2nd consistent molecular constraints
    valid = list(range(len(constraints_m)))
    if constraint2 is not None:
        for i, c in enumerate(constraints_m):
            opa1 = c[0]
            for j, opa2 in enumerate(opa_m):
                if opa2.is_conjugate(constraint2):
                    dot_m = np.dot(opa1.axis, opa2.axis)
                    # Ensure that the angles are equal
                    if abs(dot_m - dot_w) < 0.02 or abs(dot_m + dot_w) < 0.02:
                        constraints_m[i][1].append(opa2)
                        # Generate 2nd constraint in opposite direction
                        extra = deepcopy(opa2)
                        extra.axis = [
                            opa2.axis[0] * -1,
                            opa2.axis[1] * -1,
                            opa2.axis[2] * -1,
                        ]
                        constraints_m[i][1].append(extra)
            # If no consistent constraints are found, remove first constraint
            if constraints_m[i][1] == []:
                valid.remove(i)
    copy = deepcopy(constraints_m)
    constraints_m = []
    for i in valid:
        constraints_m.append(copy[i])

    # Generate orientations consistent with the possible constraints
    orientations = []
    # Loop over molecular constraint sets
    for c1 in constraints_m:
        v1 = c1[0].axis
        v2 = constraint1.axis
        T = rotate_vector(v1, v2)
        # If there is only one constraint
        if c1[1] == []:
            o = Orientation(T, degrees=1, axis=constraint1.axis)
            orientations.append(o)
        else:
            # Loop over second molecular constraints
            for opa in c1[1]:
                phi = angle(constraint1.axis, constraint2.axis)
                phi2 = angle(constraint1.axis, np.dot(T, opa.axis))
                if np.isclose(phi, phi2, rtol=rtol):
                    r = np.sin(phi)
                    c = np.linalg.norm(np.dot(T, opa.axis) - constraint2.axis)
                    theta = np.arccos(1 - (c ** 2) / (2 * (r ** 2)))
                    # R = aa2matrix(constraint1.axis, theta)
                    R = Rotation.from_rotvec(theta * constraint1.axis).as_matrix()
                    T2 = np.dot(R, T)
                    a = angle(np.dot(T2, opa.axis), constraint2.axis)
                    if not np.isclose(a, 0, rtol=rtol):
                        T2 = np.dot(np.linalg.inv(R), T)
                    o = Orientation(T2, degrees=0)
                    orientations.append(o)

    # Ensure the identity orientation is checked if no constraints are found
    if constraints_m == []:
        o = Orientation(np.identity(3), degrees=2)
        orientations.append(o)

    # Remove redundancy from orientations
    list_i = list(range(len(orientations)))
    list_j = list(range(len(orientations)))
    for i, o1 in enumerate(orientations):
        if i in list_i:
            for j, o2 in enumerate(orientations):
                if i > j and j in list_j and j in list_i:
                    # m1 = o1.get_matrix(angle=0)
                    # m2 = o2.get_matrix(angle=0)
                    m1 = o1.matrix
                    m2 = o2.matrix
                    new_op = SymmOp.from_rotation_and_translation(
                        np.dot(m2, np.linalg.inv(m1)), [0, 0, 0]
                    )
                    P = SymmOp.from_rotation_and_translation(np.linalg.inv(m1), [0, 0, 0])
                    old_op = P * new_op * P.inverse
                    if pga.is_valid_op(old_op):
                        list_i.remove(j)
                        list_j.remove(j)
    copy = deepcopy(orientations)
    orientations = []
    for i in list_i:
        orientations.append(copy[i])

    #Check each of the found orientations for consistency with the Wyckoff pos.
    #If consistent, put into an array of valid orientations
    allowed = []
    for o in orientations:
        if randomize is True:
            op = o.get_op("random")
        elif randomize is False:
            op = o.get_op() #do not change
        mo = deepcopy(mol)
        mo.apply_operation(op)
        if orientation_in_wyckoff_position(mo, wyckoff_position, exact_orientation=True, randomize=False, allow_inversion=allow_inversion) is True:
            allowed.append(o)
    #Return the array of allowed orientations. If there are none, return False
    if allowed == []:
        return False
    else:
        return allowed

def make_graph(mol, tol=0.2):
    """
    make graph object for the input molecule
    """   
    G = nx.Graph()
    names = {}
    for i, site in enumerate(mol._sites):
        names[i] = site.specie.value

    for i in range(len(mol)-1):
        site1 = mol.sites[i]
        for j in range(i+1, len(mol)):
            site2 = mol.sites[j]
            #remove short X-H distances
            if "H" in [names[i], names[j]]:
                factor = 0.3
            else:
                factor = 1.0

            if CovalentBond.is_bonded(site1, site2, factor*tol):
                G.add_edge(i,j)
    nx.set_node_attributes(G, names, 'name')

    return G
 
def compare_mol_connectivity(mol1, mol2, ignore_name=False):
    """
    Compare two molecules by connectivity
    """

    G1 = make_graph(mol1)
    G2 = make_graph(mol2)
    if ignore_name:
        GM = nx.isomorphism.GraphMatcher(G1, G2)
    else:
        fun = lambda n1, n2: n1['name'] == n2['name']
        GM = nx.isomorphism.GraphMatcher(G1, G2, node_match=fun)

    return GM.is_isomorphic(), GM.mapping


