import numpy as np
from dataclasses import dataclass
from copy import deepcopy

global_tol = 1e-8


@dataclass
class SEA():
    """
    Symmetry equivalent atoms (SEA).

    SEAs are atoms that can be swapped with no distinguishable change in the molecule.

    Parameters
    ----------
    label : str or None, optional
        Rotor type of the SEA set (e.g. Single Atom, Linear, Spherical, Regular Polygon, Oblate Symmetric Top).
    subset : np.array
        Atom indices in the molecule that belong to this SEA set, shape (N,).
    axis : np.array or None, optional
        Candidate rotational symmetry axis, shape (3,).
    """
    label: str
    subset: np.array
    axis: np.array


def transform(mol, M):
    """
    Transform coordinates of molecule by matrix M and return new molecule.

    Parameters
    ----------
    mol: ase.Atoms
        Molecule object.
    M: np.array
        Transformation matrix (e.g. rotation, reflection, etc.), shape (3,3)

    Returns
    -------
    ase.Atoms
        Molecule with transformed atom coordinates
    """
    new_mol = deepcopy(mol)
    new_mol.positions = np.dot(new_mol.positions, np.transpose(M))
    return new_mol


def distance_matrix(mol):
    """
    Calculate the interatomic distance matrix as all pairwise distances between atoms.

    Parameters
    ----------
    mol: ase.Atoms
        Molecule object.

    Returns
    -------
    np.array
        Interatomic distance matrix, shape(len(mol),len(mol))
    """
    dm = np.zeros((len(mol), len(mol)))
    for i in range(len(mol)):
        for j in range(i, len(mol)):
            dm[i,j] = np.sqrt(sum((mol.positions[i,:]-mol.positions[j,:])**2))
            dm[j,i] = dm[i,j]
    return dm


def find_SEAs(mol):
    """
    Find sets of symmetry equivalent atoms.
    Permutations of the distance matrix reveal which atoms form symmetry equivalent sets.

    Parameters
    ----------
    mol: ase.Atoms
        Molecule object.

    Returns
    -------
    List[SEA]
        List of symmetry equivalent atom sets
    """
    dm = distance_matrix(mol)
    out = []
    for i in range(len(mol)):
        for j in range(i+1, len(mol)):
            a_idx = np.argsort(dm[i,:])
            b_idx = np.argsort(dm[j,:])
            z = dm[i,a_idx] - dm[j,b_idx]
            chk = True
            for k in z:
                if abs(k) < mol.info["tol"]:
                    continue
                else:
                    chk = False
            if chk:
                out.append((i,j))
    skip = []
    SEAs = []
    for i in range(len(mol)):
        if i in skip:
            continue
        else:
            collect = [i]
        for k in out:
            if i in k:
                if i == k[0]:
                    collect.append(k[1])
                    skip.append(k[1])
                else:
                    collect.append(k[0])
                    skip.append(k[0])
        SEAs.append(SEA("", collect, np.zeros(3)))
    return SEAs


def isequivalent(A, B):
    """
    Return True if molecules A and B are equivalent with respect to permutation of like atoms.

    Parameters
    ----------
    A: ase.Atoms
        Molecule A.
    B: ase.Atoms
        Molecule B.

    Returns
    -------
    bool
        True if equivalent, False if not.
    """
    if A.info["tol"] >= B.info["tol"]:
        eq_tol = A.info["tol"]
    else:
        eq_tol = B.info["tol"]
    matched_already = []
    for i in range(len(A)):
        for j in range(len(B)):
            if j not in matched_already:
                if A.get_masses()[i] == B.get_masses()[j]:
                    zs = abs(A.positions[i,:]-B.positions[j,:])
                    if np.allclose(zs, [0,0,0], atol=eq_tol):
                        matched_already.append(j)
                        break
    if len(matched_already) == len(A):
        return True
    return False


def calcmoit(atoms):
    """
    Calculate the moment of inertia tensor for a molecule.

    Parameters
    ----------
    atoms: ase.Atoms
        Set of atoms.

    Returns
    -------
    np.array
        Cartesian moment of inertia tensor, shape(3,3).
    """
    I = np.zeros((3, 3))
    atoms.translate(-atoms.get_center_of_mass())
    masses = atoms.get_masses()
    positions = atoms.positions
    for i in range(3):
        for j in range(3):
            if i == j:
                for k in range(len(atoms)):
                    I[i,i] += masses[k]*(positions[k,(i+1)%3]**2+positions[k,(i+2)%3]**2)
            else:
                for k in range(len(atoms)):
                    I[i,j] -= masses[k]*positions[k,i]*positions[k,j]
    return I
