import numpy as np
from .sym_ops import rotation_matrix, inversion_matrix, Sn, Cn, normalize, issame_axis, isfactor, reflection_matrix
from .mol_ops import isequivalent, calcmoit, transform, find_SEAs
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import List
    from .mol_ops import SEA
    from ase import Atoms


class RotationElement():
    """
    Data structure holding a candidate rotation axis and its order.
    """
    def __init__(self, axis, order):
        self.axis = axis
        self.order = order

    def __eq__(self, other):
        if isinstance(other, RotationElement):
            return issame_axis(self.axis, other.axis) and self.order == other.order


def _intersect(a: "List[RotationElement]", b: "List[RotationElement]"):
    """Return elements present in both lists (equality by axis and order)."""
    out = []
    for re_a in a:
        for re_b in b:
            if re_a == re_b:
                out.append(re_a)
                break
    return out


def _rotation_set_intersection(rotation_set: "List[List[RotationElement]]"):
    """Return the intersection of all per-SEA rotation sets."""
    out = rotation_set[0]
    if len(rotation_set) > 1:
        for i in range(len(rotation_set)):
            out = _intersect(out, rotation_set[i])
    return out


def find_rotation_sets(mol: "Atoms", SEAs: "List[SEA]"):
    """
    For each SEA, find the set of possible RotationElements.

    Parameters
    ----------
    mol: ase.Atoms
    SEAs: List[SEA]

    Returns
    -------
    List[List[RotationElement]]
    """
    out_all_SEAs = []
    for sea in SEAs:
        length = len(sea.subset)
        out_per_SEA = []
        if length < 2:
            sea.label = "Single Atom"
        elif length == 2:
            sea.label = "Linear"
            sea.axis = normalize(mol[sea.subset[0]].position)
        else:
            sea_mol = mol[sea.subset]
            sea_mol.translate(-sea_mol.get_center_of_mass())
            evals, evecs = np.linalg.eigh(calcmoit(mol[sea.subset]))
            idx = evals.argsort()
            Ia, Ib, Ic = evals[idx]
            Iav, Ibv, Icv = [evecs[:,i] for i in idx]
            if np.isclose(Ia, Ib, atol=mol.info["tol"]) and np.isclose(Ia, Ic, atol=mol.info["tol"]):
                sea.label = "Spherical"
            elif np.isclose(Ia+Ib, Ic, atol=mol.info["tol"]):
                axis = Icv
                sea.axis = axis
                if np.isclose(Ia, Ib, atol=mol.info["tol"]):
                    sea.label = "Regular Polygon"
                    for i in range(2, length+1):
                        if isfactor(length, i):
                            out_per_SEA.append(RotationElement(axis, i))
                else:
                    sea.label = "Irregular Polygon"
                    for i in range(2, length):
                        if isfactor(length, i):
                            out_per_SEA.append(RotationElement(axis, i))
            else:
                if not (np.isclose(Ia, Ib, atol=mol.info["tol"]) or np.isclose(Ib, Ic, atol=mol.info["tol"])):
                    sea.label = "Asymmetric Rotor"
                    for ax in [Iav, Ibv, Icv]:
                        out_per_SEA.append(RotationElement(ax, 2))
                else:
                    if np.isclose(Ia, Ib, atol=mol.info["tol"]):
                        sea.label = "Oblate Symmetric Top"
                        axis = Icv
                        sea.axis = Icv
                    else:
                        sea.label = "Prolate Symmetric Top"
                        axis = Iav
                        sea.axis = Iav
                    k = length//2
                    for i in range(2, k+1):
                        if isfactor(k, i):
                            out_per_SEA.append(RotationElement(axis, i))
            if len(out_per_SEA) > 0:
                out_all_SEAs.append(out_per_SEA)
    return out_all_SEAs


def find_rotations(mol: "Atoms", rotation_set: "List[List[RotationElement]]"):
    """
    Find RotationElements in rotation_set that leave the molecule invariant.

    Parameters
    ----------
    mol: ase.Atoms
    rotation_set: List[List[RotationElement]]

    Returns
    -------
    List[RotationElement]
    """
    if len(rotation_set) < 1:
        return []
    molmoit = calcmoit(mol)
    evals = np.sort(np.linalg.eigh(molmoit)[0])
    if evals[0] == 0.0 and np.isclose(evals[1], evals[2], atol=mol.info["tol"]):
        for i in range(np.shape(mol.positions)[0]):
            if normalize(mol.positions[i,:]) is not None:
                axis = normalize(mol.positions[0,:])
        re = RotationElement(axis, 0)
        return [re]
    rsi = _rotation_set_intersection(rotation_set)
    out = []
    for i in rsi:
        rmat = Cn(i.axis, i.order)
        molB = transform(mol, rmat)
        if isequivalent(mol, molB):
            out.append(i)
    return out


def linear_mol_axis(mol: "Atoms"):
    """
    Returns the axis that best aligns with a linear molecule.

    Parameters
    ----------
    mol: ase.Atoms

    Returns
    -------
    np.array of shape (3,)
    """
    coords = mol.positions - mol.positions.mean(axis=0)
    _, _, vh = np.linalg.svd(coords, full_matrices=False)
    return normalize(vh[0])


def find_a_c2(mol: "Atoms", SEAs: "List[SEA]"):
    """
    Search for any C_2 axis; return the first one found.

    Parameters
    ----------
    mol: ase.Atoms
    SEAs: List[SEA]

    Returns
    -------
    np.array of shape (3,) or None
    """
    for sea in SEAs:
        a = _c2a(mol, sea)
        if a is not None:
            return a
        b = _c2b(mol, sea)
        if b is not None:
            return b
        if sea.label == "Linear":
            for sea2 in SEAs:
                if sea == sea2:
                    continue
                if sea2.label == "Linear":
                    c = _c2c(mol, sea, sea2)
                    if c is not None:
                        return c
    return None


def _compute_R_max(mol: "Atoms", axis: "np.array"):
    """Return the maximum perpendicular distance of any atom from the given axis."""
    axis = normalize(axis)
    coords = mol.positions
    proj = np.dot(coords, axis)[:, np.newaxis] * axis[np.newaxis, :]
    dists = np.linalg.norm(coords - proj, axis=1)
    return np.max(dists)


def is_there_ortho_c2(mol: "Atoms", SEAs: "List[SEA]", paxis: "np.array"):
    """
    Search for a C_2 axis orthogonal to paxis; return the first one found.

    Parameters
    ----------
    mol: ase.Atoms
    SEAs: List[SEA]
    paxis: np.array of shape (3,)

    Returns
    -------
    tuple(bool, np.array or None)
    """
    ortho_tol = mol.info["tol"] / _compute_R_max(mol, paxis) * 1.10
    for sea in SEAs:
        b = _c2b(mol, sea, axis=paxis)
        if b is not None and abs(np.dot(b, paxis)) <= ortho_tol:
            return True, b
        a = _c2a(mol, sea, axis=paxis)
        if a is not None and abs(np.dot(a, paxis)) <= ortho_tol:
            return True, a
        if sea.label == "Linear":
            for sea2 in SEAs:
                if sea == sea2:
                    continue
                if sea2.label == "Linear":
                    c = _c2c(mol, sea, sea2, axis=paxis)
                    if c is not None and abs(np.dot(c, paxis)) <= ortho_tol:
                        return True, c
    return False, None


def num_C2(mol: "Atoms", SEAs: "List[SEA]"):
    """
    Count distinct C_2 axes and return them.

    Parameters
    ----------
    mol: ase.Atoms
    SEAs: List[SEA]

    Returns
    -------
    tuple(int, List[np.array]) or None
    """
    axes = []
    for sea in SEAs:
        a = _c2a(mol, sea, all=True)
        if a is not None:
            axes.extend(a)
        b = _c2b(mol, sea, all=True)
        if b is not None:
            axes.extend(b)
    if len(axes) < 1:
        return None
    unique_axes = [axes[0]]
    for i in axes:
        if not any(issame_axis(i, j) for j in unique_axes):
            unique_axes.append(i)
    return len(unique_axes), unique_axes


def _c2a(mol: "Atoms", sea: "SEA", axis=None, all: bool = False):
    """Find C_2 axes from origin-midpoint vectors of atom pairs."""
    length = len(sea.subset)
    out = []
    for i in range(length):
        for j in range(i+1, length):
            midpoint = mol.positions[sea.subset[i],:] + mol.positions[sea.subset[j],:]
            if np.isclose(midpoint, [0,0,0], atol=mol.info["tol"]).all():
                continue
            midpoint = normalize(midpoint)
            if midpoint is None:
                continue
            if axis is not None and issame_axis(midpoint, axis):
                continue
            molB = transform(mol, Cn(midpoint, 2))
            if isequivalent(mol, molB):
                if all:
                    out.append(midpoint)
                else:
                    return midpoint
    return out if (all and len(out) > 0) else None


def _c2b(mol: "Atoms", sea: "SEA", axis=None, all: bool = False):
    """Find C_2 axes from individual atom position vectors."""
    length = len(sea.subset)
    out = []
    for i in range(length):
        c2_axis = normalize(mol.positions[sea.subset[i],:])
        if c2_axis is None:
            continue
        if axis is not None and issame_axis(c2_axis, axis):
            continue
        molB = transform(mol, Cn(c2_axis, 2))
        if isequivalent(mol, molB):
            if all:
                out.append(c2_axis)
            else:
                return c2_axis
    return out if (all and len(out) > 0) else None


def _c2c(mol: "Atoms", sea1: "SEA", sea2: "SEA", axis=None):
    """Find C_2 axes orthogonal to pairs of linear SEAs."""
    rij = mol.positions[sea1.subset[0],:] - mol.positions[sea1.subset[1],:]
    rkl = mol.positions[sea2.subset[0],:] - mol.positions[sea2.subset[1],:]
    c2_axis = normalize(np.cross(rij, rkl))
    if c2_axis is None:
        return None
    if axis is not None and issame_axis(c2_axis, axis):
        return None
    molB = transform(mol, Cn(c2_axis, 2))
    if isequivalent(mol, molB):
        return c2_axis
    return None


def highest_order_axis(rotations: "List[RotationElement]"):
    """Return the order of the highest-order rotation in the list."""
    return np.sort([r.order for r in rotations])[-1]


def is_there_sigmah(mol: "Atoms", paxis: "np.array"):
    """
    Check for a horizontal reflection plane (normal = paxis).

    Parameters
    ----------
    mol: ase.Atoms
    paxis: np.array of shape (3,)

    Returns
    -------
    bool
    """
    molB = transform(mol, reflection_matrix(paxis))
    return isequivalent(mol, molB)


def is_there_sigmav(mol: "Atoms", SEAs: "List[SEA]", paxis: "np.array"):
    """
    Check for vertical reflection planes (normal perpendicular to paxis).

    Parameters
    ----------
    mol: ase.Atoms
    SEAs: List[SEA]
    paxis: np.array of shape (3,)

    Returns
    -------
    tuple(bool, np.array or None)
    """
    axes = []
    for sea in SEAs:
        length = len(sea.subset)
        if length < 2:
            continue
        A = sea.subset[0]
        for i in range(1, length):
            B = sea.subset[i]
            n = normalize(mol.positions[A,:] - mol.positions[B,:])
            if n is not None:
                molB = transform(mol, reflection_matrix(n))
                if isequivalent(mol, molB):
                    axes.append(n)
    if len(axes) < 1:
        if mol_is_planar(mol):
            return True, planar_mol_axis(mol)
        else:
            return False, None
    unique_axes = [axes[0]]
    for i in axes:
        if not any(issame_axis(i, j) for j in unique_axes):
            unique_axes.append(i)
    for i in unique_axes:
        if not issame_axis(i, paxis):
            return True, i
    return False, None


def mol_is_planar(mol: "Atoms"):
    """
    Check if all atoms lie in a common plane.

    Parameters
    ----------
    mol: ase.Atoms

    Returns
    -------
    bool
    """
    rank = np.linalg.matrix_rank(mol.positions, tol=mol.info["tol"])
    if rank < 3:
        from .symtext import rotate_mol_to_symels
        axis = planar_mol_axis(mol)
        new_mol, _, _ = rotate_mol_to_symels(mol, axis, np.array([0, 0, 0]))
        matrix = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]])
        molB = transform(new_mol, matrix)
        for i in range(len(mol)):
            if not np.isclose(new_mol.positions[i,:], molB.positions[i,:], atol=mol.info["tol"]).all():
                return False
        return True
    return False


def planar_mol_axis(mol: "Atoms"):
    """
    Return the normal to the plane of a planar molecule.

    Parameters
    ----------
    mol: ase.Atoms

    Returns
    -------
    np.array of shape (3,) or None
    """
    coords = mol.positions - mol.positions.mean(axis=0)
    _, _, vh = np.linalg.svd(coords, full_matrices=False)
    return normalize(vh[-1])


def find_C3s_for_Ih(mol: "Atoms"):
    """
    Find the ten unique C_3 axes of an Ih molecule.

    Parameters
    ----------
    mol: ase.Atoms

    Returns
    -------
    List[np.array of shape (3,)]
    """
    c3_axes = []
    for i in range(len(mol)):
        for j in range(len(mol)):
            for k in range(len(mol)):
                if i != j and i != k:
                    rij = mol.positions[i,:] - mol.positions[j,:]
                    rjk = mol.positions[j,:] - mol.positions[k,:]
                    rik = mol.positions[i,:] - mol.positions[k,:]
                    nij, njk, nik = np.linalg.norm(rij), np.linalg.norm(rjk), np.linalg.norm(rik)
                    if np.isclose(nij, njk, atol=mol.info["tol"]) and np.isclose(nij, nik, atol=mol.info["tol"]):
                        c3_axis = normalize(np.cross(rij, rjk))
                        if c3_axis is not None:
                            molB = transform(mol, Cn(c3_axis, 3))
                            if isequivalent(mol, molB):
                                c3_axes.append(c3_axis)
    unique_axes = [c3_axes[0]]
    for i in c3_axes:
        if not any(issame_axis(i, j) for j in unique_axes):
            unique_axes.append(i)
    chk = len(unique_axes)
    if chk != 10:
        raise Exception(f"Unexpected number of C3 axes for Ih: found {chk}, expected 10")
    return unique_axes


def find_C4s_for_Oh(mol: "Atoms"):
    """
    Find the three unique C_4 axes of an Oh molecule.

    Parameters
    ----------
    mol: ase.Atoms

    Returns
    -------
    List[np.array of shape (3,)]
    """
    c4_axes = []
    for i in range(len(mol)):
        for j in range(len(mol)):
            for k in range(len(mol)):
                for l in range(len(mol)):
                    if i != j and k != l and i != k:
                        rij = mol.positions[i,:] - mol.positions[j,:]
                        rjk = mol.positions[j,:] - mol.positions[k,:]
                        rkl = mol.positions[k,:] - mol.positions[l,:]
                        ril = mol.positions[i,:] - mol.positions[l,:]
                        nij, njk = np.linalg.norm(rij), np.linalg.norm(rjk)
                        nkl, nil_ = np.linalg.norm(rkl), np.linalg.norm(ril)
                        if (np.isclose(nij, njk, atol=mol.info["tol"]) and
                                np.isclose(nkl, nil_, atol=mol.info["tol"]) and
                                np.isclose(nij, nkl, atol=mol.info["tol"])):
                            c4_axis = normalize(np.cross(rij, rjk))
                            if c4_axis is not None:
                                molB = transform(mol, Cn(c4_axis, 4))
                                if isequivalent(mol, molB):
                                    c4_axes.append(c4_axis)
    unique_axes = [c4_axes[0]]
    for i in c4_axes:
        if not any(issame_axis(i, j) for j in unique_axes):
            unique_axes.append(i)
    chk = len(unique_axes)
    if chk != 3:
        raise Exception(f"Unexpected number of C4 axes for Oh: found {chk}, expected 3")
    return unique_axes


def find_point_group(mol: "Atoms"):
    """
    Find the point group of a molecule.

    Returns the Schoenflies symbol and the primary/secondary axes that define
    the canonical orientation with respect to the generated symmetry elements.

    Based on:
        Beruski, Otávio; Vidal, Luciano N. J. Comp. Chem. 2013. doi:10.1002/jcc.23493

    Parameters
    ----------
    mol: ase.Atoms

    Returns
    -------
    tuple(str, (np.array, np.array))
        Point group string, primary axis, secondary axis; axes have shape (3,).
    """
    try:
        mol.info['tol']
    except KeyError:
        raise Exception("Atoms object tolerance hasn't been set. Set it with Atoms.info['tol']=.")

    mol = mol.copy()
    paxis = [0,0,0]
    saxis = [0,0,0]
    moit = calcmoit(mol)
    Ia_mol, Ib_mol, Ic_mol = np.sort(np.linalg.eigh(moit)[0])

    # Linear tops
    if np.isclose(Ia_mol, 0.0, atol=mol.info["tol"]):
        paxis = linear_mol_axis(mol)
        pg = "D0h" if isequivalent(mol, transform(mol, inversion_matrix())) else "C0v"

    # Spherical tops
    elif np.isclose(Ia_mol, Ib_mol, atol=mol.info["tol"]) and np.isclose(Ia_mol, Ic_mol, atol=mol.info["tol"]):
        seas = find_SEAs(mol)
        n, axes = num_C2(mol, seas)
        invertable = isequivalent(mol, transform(mol, inversion_matrix()))
        if n == 15:  # Icosahedral
            tempaxis = axes[0]
            c3s = find_C3s_for_Ih(mol)
            for c3 in c3s:
                if np.isclose(np.arccos(abs(np.dot(c3, tempaxis))), 0.36486382647383764, atol=1e-4):
                    taxis = normalize(np.cross(c3, tempaxis))
                    saxis = normalize(np.cross(taxis, tempaxis))
                    break
            phi = (1+np.sqrt(5.0))/2
            theta = np.arccos(phi/np.sqrt(1+(phi**2)))
            rmat = rotation_matrix(saxis, theta)
            paxis = np.dot(rmat, tempaxis)
            pg = "Ih" if invertable else "I"
        elif n == 9:  # Octahedral
            c4s = find_C4s_for_Oh(mol)
            paxis = c4s[0]
            saxis = c4s[1]
            pg = "Oh" if invertable else "O"
        elif n == 3:  # Tetrahedral
            paxis = axes[0]
            saxis = axes[1]
            pg = "Th" if invertable else "Td"

    # Asymmetric and symmetric tops
    else:
        seas = find_SEAs(mol)
        rot_set = find_rotation_sets(mol, seas)
        rots = find_rotations(mol, rot_set)
        if len(rots) >= 1:
            Cn_order = highest_order_axis(rots)
            paxis = rots[0].axis
        else:
            c2 = find_a_c2(mol, seas)
            if c2 is None:
                if isequivalent(mol, transform(mol, inversion_matrix())):
                    return "Ci", (paxis, saxis)
                sigmav_chk, sigmav = is_there_sigmav(mol, seas, np.asarray([0,0,0]))
                if sigmav_chk:
                    if sigmav is not None:
                        paxis = sigmav
                    return "Cs", (paxis, saxis)
                return "C1", (paxis, saxis)
            paxis = c2
            Cn_order = 2

        ortho_c2_chk, c2_ortho = is_there_ortho_c2(mol, seas, paxis)
        sigmav_chk, sigmav = is_there_sigmav(mol, seas, paxis)
        sigmah_chk = is_there_sigmah(mol, paxis)
        if ortho_c2_chk:
            saxis = c2_ortho
            if sigmah_chk:
                pg = "D"+str(Cn_order)+"h"
            elif sigmav_chk:
                pg = "D"+str(Cn_order)+"d"
            else:
                pg = "D"+str(Cn_order)
        elif sigmah_chk:
            pg = "C"+str(Cn_order)+"h"
        elif sigmav_chk:
            pg = "C"+str(Cn_order)+"v"
            if mol_is_planar(mol):
                saxis = planar_mol_axis(mol)
            elif sigmav.any():
                saxis = normalize(np.cross(paxis, sigmav))
        else:
            S2n = Sn(paxis, Cn_order*2)
            molB = transform(mol, S2n)
            pg = "S"+str(2*Cn_order) if isequivalent(mol, molB) else "C"+str(Cn_order)

    return pg, (paxis, saxis)
