from pymatgen.io.ase import AseAtomsAdaptor
from pymatgen.core import Molecule
from pymatgen.symmetry.analyzer import PointGroupAnalyzer
#------------------------------------------------------------------------------------------
def molase2pymatgen(moleculease):
    """Convert an ASE Atoms object to the corresponding pymatgen object.

in:
    moleculease (ase.Atoms): structure to convert. Periodic structures (pbc=True)
        are converted to a pymatgen Structure; isolated molecules to a pymatgen Molecule.
out:
    pymatgen.core.Structure or pymatgen.core.Molecule: the converted object.
"""
    if any(moleculease.pbc):
        structurepymatgen = AseAtomsAdaptor.get_structure(moleculease)
        return structurepymatgen
    else:
        symbols = moleculease.get_chemical_symbols()
        positions = moleculease.get_positions()
        moleculepymatgen=Molecule(symbols, positions)
        return moleculepymatgen
#------------------------------------------------------------------------------------------
def point_group(moleculease, tolerance=0.25, eigen_tolerance=1E-3, matrix_tolerance=0.1):
    """Determine the point group of a molecule using pymatgen's PointGroupAnalyzer.

in:
    moleculease (ase.Atoms): molecular structure (non-periodic).
    tolerance (float): distance tolerance in Angstroms for symmetry detection (default 0.25).
    eigen_tolerance (float): tolerance for comparing eigenvalues (default 1e-3).
    matrix_tolerance (float): tolerance for comparing rotation matrices (default 0.1).
out:
    pymatgen.symmetry.analyzer.PointGroup: the identified point group object.
"""
    molpmg=molase2pymatgen(moleculease)
    molx=PointGroupAnalyzer(molpmg, tolerance, eigen_tolerance, matrix_tolerance)
    pointgroup=molx.get_pointgroup()
    return pointgroup
#------------------------------------------------------------------------------------------
def symmetry_number(moleculease, tolerance=0.25, eigen_tolerance=1E-3, matrix_tolerance=0.1):
    """Return the number of symmetry operations (symmetry number) of a molecule.

in:
    moleculease (ase.Atoms): molecular structure (non-periodic).
    tolerance (float): distance tolerance in Angstroms (default 0.25).
    eigen_tolerance (float): eigenvalue comparison tolerance (default 1e-3).
    matrix_tolerance (float): rotation matrix comparison tolerance (default 0.1).
out:
    int: total number of symmetry operations found for the molecule.
"""
    molpmg=molase2pymatgen(moleculease)
    molx=PointGroupAnalyzer(molpmg, tolerance, eigen_tolerance, matrix_tolerance)
    symmetry_number=len(molx.get_symmetry_operations())
    return symmetry_number
#------------------------------------------------------------------------------------------
def inequivalent_finder(moleculease, tolerance=0.3, eigen_tolerance=0.008, matrix_tolerance=0.1):
    """Find the indices of symmetry-inequivalent atoms in a molecule.

in:
    moleculease (ase.Atoms): molecular structure (non-periodic).
    tolerance (float): distance tolerance in Angstroms for symmetry detection (default 0.3).
    eigen_tolerance (float): eigenvalue comparison tolerance (default 0.008).
    matrix_tolerance (float): rotation matrix comparison tolerance (default 0.1).
out:
    list of int: sorted list of atom indices representing one atom from each
    symmetry-equivalent set (i.e., the inequivalent sites).
"""
    molpmg=molase2pymatgen(moleculease)
    #molpmg=AseAtomsAdaptor.get_structure(moleculease)
    molx=PointGroupAnalyzer(molpmg, tolerance, eigen_tolerance, matrix_tolerance)
    dictionary=molx.get_equivalent_atoms()
    dict=dictionary['eq_sets']
    inequiv=[ix for ix in dict.keys()]
    #for ix in dict.keys(): print(list(dict[ix]))
    #print('')
    inequiv.sort()
    return inequiv
#------------------------------------------------------------------------------------------
