import numpy as np
from .pg_detect import find_point_group
from .point_group import PointGroup
from .symel_gen import pg_to_symels
from .symel import Symel
from .mol_ops import transform, global_tol
from .sym_ops import normalize


def rotate_mol_to_symels(mol, paxis, saxis):
    """
    Rotate molecule so that paxis aligns with z and saxis aligns with x.

    Returns the rotated molecule and the forward/inverse rotation matrices so
    that computed properties can be rotated back to the original orientation.

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

    Returns
    -------
    tuple(ase.Atoms, np.array, np.array)
        Rotated molecule, rotation matrix, inverse rotation matrix; matrices shape (3,3).
    """
    if np.isclose(np.linalg.norm(paxis), 0.0, atol=global_tol):
        rmat = rmat_inv = np.eye(3)
        return mol, rmat, rmat_inv
    z = paxis
    if np.isclose(np.linalg.norm(saxis), 0.0, atol=global_tol):
        x = None
        for trial_vec in np.eye(3):
            x = np.cross(trial_vec, z)
            if not np.isclose(np.linalg.norm(x), 0, atol=global_tol):
                x = normalize(x)
                break
        y = normalize(np.cross(z, x))
    else:
        x = saxis
        y = np.cross(z, x)
    rmat = np.column_stack((x, y, z))
    rmat_inv = rmat.T
    new_mol = transform(mol, rmat_inv)
    return new_mol, rmat, rmat_inv


def _where_you_go(mol, atom, symel):
    """
    Find the atom index that atom maps to under symel.

    Parameters
    ----------
    mol: ase.Atoms
    atom: int
    symel: Symel

    Returns
    -------
    int or None
    """
    ratom = np.dot(symel.rrep, mol.positions[atom,:].T)
    for i in range(len(mol)):
        if np.isclose(mol.positions[i,:], ratom, atol=mol.info["tol"]).all():
            return i
    return None


def _get_atom_mapping(mol, symels):
    """
    Build the (natom × nsymel) atom permutation map.

    Parameters
    ----------
    mol: ase.Atoms
    symels: List[Symel]

    Returns
    -------
    np.array of shape (natom, nsymel)
    """
    amap = np.zeros((len(mol), len(symels)), dtype=int)
    for atom in range(len(mol)):
        for s, symel in enumerate(symels):
            w = _where_you_go(mol, atom, symel)
            if w is not None:
                amap[atom, s] = w
            else:
                raise Exception(
                    f"Atom {atom} {mol.info['num']} not mapped to another atom "
                    f"under symel {symel}\nPositions: {mol.positions}"
                )
    return amap


def _get_linear_atom_mapping(mol, pg):
    """
    Atom permutation map for linear point groups.

    Parameters
    ----------
    mol: ase.Atoms
    pg: PointGroup

    Returns
    -------
    np.array
    """
    amap = np.array([atom for atom in range(len(mol))], dtype=int).reshape((len(mol), 1))
    if pg.family == "D":
        ungerade_map = np.zeros((len(mol)), dtype=int)
        for atom in range(len(mol)):
            w = _where_you_go(mol, atom, Symel("i", None, -1*np.eye(3), None, None, None))
            if w is not None:
                ungerade_map[atom] = w
            else:
                raise Exception(f"Atom {atom} not mapped to another atom under symel i")
        return np.column_stack((amap, ungerade_map))
    return amap


class Symtext():
    """
    Full symmetry characterization of a molecule.

    Holds the point group, list of symmetry elements, atom permutation map,
    and the rotation matrices that relate the molecule to the canonical
    orientation defined by the symmetry elements.
    """
    def __init__(self, mol, rotate_to_std, reverse_rotate, pg, symels, atom_map) -> None:
        self.mol = mol
        self.rotate_to_std = rotate_to_std
        self.reverse_rotate = reverse_rotate
        self.pg = pg
        self.symels = symels
        self.atom_map = atom_map
        if pg.is_linear:
            self.order = 4*np.pi if pg.family == "C" else 8*np.pi
        else:
            self.order = len(symels)

    def __len__(self):
        return len(self.symels)

    def __repr__(self):
        return f"\n{self.mol}\n{self.symels}\nAtom map:\n{self.atom_map}"

    @classmethod
    def empty(cls):
        return Symtext(None, None, None, PointGroup.from_string("C1"), [], [])

    @classmethod
    def from_molecule(cls, mol):
        """
        Build a Symtext from an ase.Atoms object.

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

        Returns
        -------
        Symtext
        """
        mol.translate(-mol.get_center_of_mass())
        pg_str, (paxis, saxis) = find_point_group(mol)
        pg = PointGroup.from_string(pg_str)
        mol, reverse_rotate, rotate_to_std = rotate_mol_to_symels(mol, paxis, saxis)
        mol.info["pg"] = pg_str
        symels = pg_to_symels(pg.str)
        if pg.is_linear:
            atom_map = _get_linear_atom_mapping(mol, pg)
        else:
            atom_map = _get_atom_mapping(mol, symels)
        return Symtext(mol, rotate_to_std, reverse_rotate, pg, symels, atom_map)
