import os
import time
import numpy as np
from multiprocessing import Process
from ase.data import atomic_masses
from aegon.libutils import readxyzs, writexyzs, centroid, sort_by_energy, prepare_folders, split_poscarlist, align_two
tolsij=0.95
tolene=0.10
#------------------------------------------------------------------------------------------
#Ultrafast Shape Recognition (USR) Algorithm
#------------------------------------------------------------------------------------------
def lastwo_central_moment(listxxx):
    """Compute the second and third central moments of a distribution.

in:
    listxxx (numpy.ndarray): 1D array of distance values.
out:
    list of float: [variance, third central moment] of the input distribution.
"""
    xavg=np.mean(listxxx)
    res = listxxx - xavg
    ssq2 = np.mean(res**2)
    ssq3 = np.mean(res**3)
    return [ssq2,ssq3]
#------------------------------------------------------------------------------------------
def four_points(atoms):
    """Compute the four reference points used in the USR shape descriptor.

in:
    atoms (ase.Atoms): molecular structure.
out:
    tuple: (ctd, cst, fct, ftf) where ctd is the centroid, cst is the atom
    closest to ctd, fct is the farthest atom from ctd, and ftf is the farthest
    atom from fct. All are numpy.ndarray of shape (3,).
"""
    ####ctd: the molecular centroid
    ####cst: The closest atom to the molecular centroid (ctd)
    ####fct: The farthest atom from the molecular centroid (ctd)
    ####ftf: The farthest atom from the fct
    vecpos=[np.array(position) for position in atoms.get_positions()]
    ctd=centroid(atoms)
    r=[[np.linalg.norm(xyz - ctd), xyz] for xyz in vecpos]
    r.sort(key=lambda x: x[0])
    cst=r[0][1]
    fct=r[-1][1]
    s=[[np.linalg.norm(xyz - fct), xyz] for xyz in vecpos]
    s.sort(key=lambda x: x[0])
    ftf=s[-1][1]
    return ctd, cst, fct, ftf
#------------------------------------------------------------------------------------------
def USRMonoatom(moleculein):
    """Compute the 8-feature USR shape descriptor for a monatomic cluster.

in:
    moleculein (ase.Atoms): monatomic cluster (single element species).
out:
    list of float: 8-element USR descriptor (2 moments x 4 reference points).
"""
    ctd,   cst,  fct,  ftf=four_points(moleculein)
    lctd, lcst, lfct, lftf= [], [], [], []
    for iatom in moleculein:
        xyz=np.array(iatom.position)
        ictd1=np.linalg.norm(xyz-ctd)
        icst1=np.linalg.norm(xyz-cst)
        ifct1=np.linalg.norm(xyz-fct)
        iftf1=np.linalg.norm(xyz-ftf)
        lctd.append(ictd1)
        lcst.append(icst1)
        lfct.append(ifct1)
        lftf.append(iftf1)
    a1=lastwo_central_moment(lctd)
    a2=lastwo_central_moment(lcst)
    a3=lastwo_central_moment(lfct)
    a4=lastwo_central_moment(lftf)
    standard_USR=a1+a2+a3+a4
    return standard_USR
#------------------------------------------------------------------------------------------
def USRMultiatom(moleculein):
    """Compute the 16-feature mass-weighted USR descriptor for a multielement cluster.

in:
    moleculein (ase.Atoms): cluster which may contain multiple element types.
out:
    list of float: 16-element USR descriptor combining 8 standard and 8 mass-weighted
    distance-moment features.
"""
    listm=moleculein.get_masses()
    mavg = np.mean(listm)
    ctd, cst, fct, ftf=four_points(moleculein)
    lctd,  lcst,  lfct,  lftf = [], [], [], []
    lctdm, lcstm, lfctm, lftfm = [], [], [], []
    for iatom in moleculein:
        xyz=np.array(iatom.position)
        ictd1=np.linalg.norm(xyz-ctd)
        icst1=np.linalg.norm(xyz-cst)
        ifct1=np.linalg.norm(xyz-fct)
        iftf1=np.linalg.norm(xyz-ftf)
        lctd.append(ictd1)
        lcst.append(icst1)
        lfct.append(ifct1)
        lftf.append(iftf1)
        mi=atomic_masses[iatom.number]
        ictdm1=ictd1*mi/mavg
        icstm1=icst1*mi/mavg
        ifctm1=ifct1*mi/mavg
        iftfm1=iftf1*mi/mavg
        lctdm.append(ictdm1)
        lcstm.append(icstm1)
        lfctm.append(ifctm1)
        lftfm.append(iftfm1)
    a1=lastwo_central_moment(lctd)
    a2=lastwo_central_moment(lcst)
    a3=lastwo_central_moment(lfct)
    a4=lastwo_central_moment(lftf)
    b1=lastwo_central_moment(lctdm)
    b2=lastwo_central_moment(lcstm)
    b3=lastwo_central_moment(lfctm)
    b4=lastwo_central_moment(lftfm)
    standard_USR=a1+a2+a3+a4
    m_weight_USR=b1+b2+b3+b4
    extra_descriptor_USR=standard_USR+m_weight_USR
    return extra_descriptor_USR
#------------------------------------------------------------------------------------------
def similarity(imol, jmol, mono=False):
    """Compute the USR similarity score between two clusters.

in:
    imol (ase.Atoms): first cluster.
    jmol (ase.Atoms): second cluster.
    mono (bool): if True, use the 8-feature monatomic USR; if False (default),
        use the 16-feature multiatom USR.
out:
    float: similarity score in (0, 1]. Values close to 1 indicate similar shapes.
"""
    if mono:
        n_features=8
        vi=USRMonoatom(imol)
        vj=USRMonoatom(jmol)
    else:
        n_features=16
        vi=USRMultiatom(imol)
        vj=USRMultiatom(jmol)
    manhattan_distance=sum([np.absolute(a-b) for a,b in zip(vi,vj)])
    sij=1.0/(1.0+manhattan_distance/float(n_features))
    return sij
#------------------------------------------------------------------------------------------
def find_similar_elements(listmol, similarity_matrix, tols=tolsij, tole=tolene):
    """Filter similar structures from a list given a precomputed similarity matrix.

in:
    listmol (list of ase.Atoms): list of structures; each must have info['e'].
    similarity_matrix (numpy.ndarray): NxN matrix of pairwise USR similarities.
    tols (float): similarity threshold above which two structures are considered similar.
    tole (float): energy difference threshold (in energy units) for flagging similarity.
out:
    list of ase.Atoms: dissimilar structures (those not flagged as duplicates).
    Writes a 'debug.xyz' file of aligned pairs for visual inspection.
"""
    similar_elements_indices = []
    num_elements = similarity_matrix.shape[0]
    debug=[]
    for i in range(num_elements):
        for j in range(i+1,num_elements):
            edf=np.abs(listmol[i].info['e'] - listmol[j].info['e'])
            if (similarity_matrix[i, j] >= tols) and (edf <= tole):
                similar_elements_indices.append(j)
                x1=listmol[j].copy()
                x2=listmol[i].copy()
                y1,y2=align_two(x1, x2)
                debug.extend([y1])
                debug.extend([y2])
    similar_elements_indices.sort()
    disimilars_atoms=[listmol[i] for i in range(num_elements) if i not in similar_elements_indices]
    if len(debug) > 0: writexyzs(debug, 'debug.xyz')
    return disimilars_atoms
#------------------------------------------------------------------------------------------
def disc_USR(listmol, tols=tolsij, tole=tolene, mono=False):
    """Discriminate a list of structures by USR shape similarity and energy proximity.

in:
    listmol (list of ase.Atoms): structures to discriminate; each must have info['e'].
    tols (float): USR similarity threshold (default 0.95).
    tole (float): energy proximity threshold (default 0.10).
    mono (bool): if True, use the monatomic USR descriptor.
out:
    list of ase.Atoms: subset of listmol containing only dissimilar structures.
"""
    num_molecules = len(listmol)
    if mono:
        n_features=8
        descriptors=[USRMonoatom(imol) for imol in listmol]
    else:
        n_features=16
        descriptors=[USRMultiatom(imol) for imol in listmol]
    #similar_elements_indices = []
    similarity_matrix = np.zeros((num_molecules, num_molecules))
    for i in range(num_molecules):
        vi=descriptors[i]
        for j in range(i, num_molecules):
            vj=descriptors[j]
            manhattan_distance=sum([np.absolute(a-b) for a,b in zip(vi,vj)])
            similarity=1.0/(1.0+manhattan_distance/float(n_features))
            similarity_matrix[i, j] = similarity
            similarity_matrix[j, i] = similarity
    disimilars_atoms=find_similar_elements(listmol, similarity_matrix, tols, tole)
    return disimilars_atoms
#------------------------------------------------------------------------------------------
def comparator_usr_serial(atoms0, tols=tolsij, tole=tolene, mono=False):
    """Run serial USR discrimination on a list of structures and print timing.

in:
    atoms0 (list of ase.Atoms): structures to compare; each must have info['e'].
    tols (float): USR similarity threshold (default 0.95).
    tole (float): energy proximity threshold (default 0.10).
    mono (bool): if True, use the monatomic USR descriptor.
out:
    list of ase.Atoms: dissimilar structures after filtering.
"""
    start = time.time()
    ni=len(atoms0)
    atoms1=disc_USR(atoms0, tols, tole, mono)
    nf=len(atoms1)
    end = time.time()
    print('USR comparison (-serial-) at %5.2f s [%d -> %d]' %(end - start, ni, nf))
    return atoms1
#------------------------------------------------------------------------------------------
def make_comparator_usr(ifolder, tols=tolsij, tole=tolene, mono=False):
    """Run USR discrimination on structures stored in a folder and write the result.

in:
    ifolder (str): folder name; reads ifolder/ifolder.xyz and writes
        ifolder/ifolder_disc.xyz.
    tols (float): USR similarity threshold.
    tole (float): energy proximity threshold.
    mono (bool): if True, use the monatomic USR descriptor.
out:
    None. Writes the discriminated structures to disk.
"""
    atoms0=readxyzs(ifolder+'/'+ifolder+'.xyz')
    atoms1=disc_USR(atoms0, tols, tole, mono)
    writexyzs(atoms1,ifolder+'/'+ifolder+'_disc.xyz')
#------------------------------------------------------------------------------------------
def comparator_usr_parallel(poscarlist, nproc, base_name, tols=tolsij, tole=tolene, mono=False):
    """Run parallel USR discrimination by splitting the list across subprocesses.

in:
    poscarlist (list of ase.Atoms): structures to discriminate; each must have info['e'].
    nproc (int): number of parallel processes.
    base_name (str): base prefix for temporary folder names.
    tols (float): USR similarity threshold.
    tole (float): energy proximity threshold.
    mono (bool): if True, use the monatomic USR descriptor.
out:
    list of ase.Atoms: discriminated structures sorted by energy. Falls back to
    serial comparison if the list is smaller than 4*nproc. Prints timing.
"""
    poscarlist=sort_by_energy(poscarlist, 1)
    ni=len(poscarlist)
    if ni < 4*nproc:
        moleculeout=comparator_usr_serial(poscarlist, tols, tole, mono)
        moleculeout=sort_by_energy(moleculeout, 1)
        return moleculeout
    start = time.time()
    folderlist=prepare_folders(poscarlist, nproc, base_name)
    poscar_split_list=split_poscarlist(poscarlist, nproc)
    procs = []
    dicc_term = {iposcar.info['i']: iposcar.info['c'] for iposcar in poscarlist}
    for ifolder, iposcars in zip(folderlist, poscar_split_list):
        writexyzs(iposcars, ifolder+'/'+ifolder+'.xyz')
        proc = Process(target=make_comparator_usr, args=(ifolder, tols, tole, mono))
        procs.append(proc)
        proc.start()
    for proc in procs:
        proc.join()
    center_mols, border_mols = [], []
    for ifolder in folderlist:
        molx=readxyzs(ifolder+'/'+ifolder+'_disc.xyz')
        for imol in molx: imol.info['c']=dicc_term[imol.info['i']]
        molx=sort_by_energy(molx, 1)
        emin=molx[0].info['e']
        emax=molx[-1].info['e']
        for imol in molx:
            delw=imol.info['e'] - emin
            deup=emax - imol.info['e']
            if ( delw < tole ) or ( deup < tole ):
                border_mols.extend([imol])
            else:
                center_mols.extend([imol])
    border_mols=disc_USR(border_mols, tols, tole, mono)
    moleculeout=center_mols+border_mols
    moleculeout=sort_by_energy(moleculeout, 1)
    os.system('rm -rf %sproc[0-9][0-9][0-9]' %(base_name))
    nf=len(moleculeout)
    end = time.time()
    print('USR comparison (parallel) at %5.2f s [%d -> %d]' %(end - start, ni, nf))
    return moleculeout
#------------------------------------------------------------------------------------------
def make_similarity_matrix_compare(moleculein, moleculeref, mono=False):
    """Build a cross-similarity matrix between two lists of structures.

in:
    moleculein (list of ase.Atoms): first set of structures (N1 items).
    moleculeref (list of ase.Atoms): reference set of structures (N2 items).
    mono (bool): if True, use the monatomic 8-feature USR descriptor.
out:
    numpy.ndarray: N1 x N2 matrix of USR similarity values in (0, 1].
"""
    if mono:
        n_features=8
        descriptors1=[USRMonoatom(imol) for imol in moleculein]
        descriptors2=[USRMonoatom(imol) for imol in moleculeref]
    else:
        n_features=16
        descriptors1=[USRMultiatom(imol) for imol in moleculein]
        descriptors2=[USRMultiatom(imol) for imol in moleculeref]
    total_molecules1=len(moleculein)
    total_molecules2=len(moleculeref)
    similarity_matrix=np.zeros(shape=(total_molecules1, total_molecules2),dtype=float)
    for i in range(total_molecules1):
        vi=descriptors1[i]
        for j in range(total_molecules2):
            vj=descriptors2[j]
            manhattan_distance=sum([np.absolute(a-b) for a,b in zip(vi,vj)])
            similarity=1.0/(1.0+manhattan_distance/float(n_features))
            similarity_matrix[i, j] = similarity
    return similarity_matrix
#------------------------------------------------------------------------------------------
def molin_sim_molref(moleculein, moleculeref, tols=tolsij, tole=tolene, mono=False):
    """Remove from moleculein any structure that is similar to one in moleculeref.

in:
    moleculein (list of ase.Atoms): candidate structures; each must have info['e'].
    moleculeref (list of ase.Atoms): reference structures to compare against.
    tols (float): USR similarity threshold (default 0.95).
    tole (float): energy proximity threshold (default 0.10).
    mono (bool): if True, use the monatomic USR descriptor.
out:
    list of ase.Atoms: subset of moleculein containing structures not similar
    to any reference structure. Prints timing.
"""
    start = time.time()
    ni=len(moleculein)
    matrixs=make_similarity_matrix_compare(moleculein, moleculeref, mono)
    similares=[]
    for i,imol in enumerate(moleculein):
        for j, jmol in enumerate(moleculeref):
            sij=matrixs[i][j]
            edf=np.abs(imol.info['e']-jmol.info['e'])
            if (sij >= tols) and (edf <= tole):
                similares.append(i)
                break
    moleculeout=[imol for i, imol in enumerate(moleculein) if i not in similares]
    nf=len(moleculeout)
    end = time.time()
    print('USR comparison (-serial-) at %5.2f s [%d -> %d]' %(end - start, ni, nf))
    return moleculeout
#------------------------------------------------------------------------------------------
