import random
import numpy as np
from ase import Atoms
from aegon.libsel_roulette import get_roulette_wheel_selection
from aegon.libutils        import rand_unit_vector, rotate_random, distance_atoms2atoms, rotate_vector_angle_deg, rename
from aegon.libdiscusr      import molin_sim_molref, comparator_usr_serial
#------------------------------------------------------------------------------------------
""" CUT IN THE Z AXIS """
def get_gen(atoms):
    """Split an Atoms object into two halves at the Z=0 plane.

in:
    atoms (ase.Atoms): cluster centered near the origin.
out:
    tuple: (children_up, children_dn) where children_up contains atoms with z>=0
    and children_dn those with z<0. Each has info['c'] storing the original indices.
"""
    ene=float(0.0)
    list_up, list_dn = [], []
    children_up=Atoms()
    children_dn=Atoms()
    for ii, iatom in enumerate(atoms):
        if iatom.position[2] >= 0.0:
            children_up.append(iatom)
            list_up.append(ii)
        else:
            children_dn.append(iatom)
            list_dn.append(ii)
    children_up.info['c']=list_up
    children_dn.info['c']=list_dn
    return children_up, children_dn
#------------------------------------------------------------------------------------------
def randvector(r):
    """Generate a random displacement vector with magnitude uniformly sampled in [0, r].

in:
    r (float): maximum displacement magnitude.
out:
    numpy.ndarray: 3D displacement vector of length in [0, r] pointing in a random direction.
"""
    ru=float(random.uniform(0.0, r))
    vector = ru*rand_unit_vector()
    return vector
#------------------------------------------------------------------------------------------
def molecule_from_list(moleculein, n):
    """Build a sub-Atoms object containing only the atoms at the specified indices.

in:
    moleculein (ase.Atoms): source structure.
    n (int or list of int): single index or list of indices to extract.
out:
    ase.Atoms: new Atoms object containing the selected atoms.
"""
    if not isinstance(n, list): n = [n]
    moleculeout = Atoms()
    for i in n: moleculeout.append(moleculein[i])
    return moleculeout
#------------------------------------------------------------------------------------------
def up_part_from_n(father1, nff, deltadisp):
    """Find a random rotation/translation of father1 such that exactly nff atoms lie above Z=0.

in:
    father1 (ase.Atoms): parent cluster to cut.
    nff (int): required number of atoms in the upper (z>=0) half.
    deltadisp (float): maximum random displacement applied before rotation.
out:
    tuple: (fup, fdn) upper and lower halves when nff atoms are achieved;
    False if not achieved within 5000 attempts.
"""
    fall=True
    contador=0
    while fall:
        father2=father1.copy()
        rv=randvector(deltadisp)
        father2.translate(-rv)
        rotate_random(father2)
        fup, fdn=get_gen(father2)
        fup_cs=fup.get_chemical_symbols()
        fdn_cs=fdn.get_chemical_symbols()
        if len(fup_cs)==nff:
            fall=False
            return fup, fdn
        elif len(fdn_cs)==nff:
            fall=False
            positions = fup.get_positions()
            positions[:, 2] *= -1
            fup.set_positions(positions)
            positions = fdn.get_positions()
            positions[:, 2] *= -1
            fdn.set_positions(positions)
            return fdn, fup
        else:
            contador=contador+1
            if contador==5000:
                return False
#------------------------------------------------------------------------------------------
def dn_part_from_zlist(mother1, list_chemical_symbols, deltadisp):
    """Find a rotation/translation of mother1 so the lower half matches the target species list.

in:
    mother1 (ase.Atoms): parent cluster to cut.
    list_chemical_symbols (list of str): required chemical symbols for the lower half.
    deltadisp (float): maximum random displacement applied before rotation.
out:
    tuple: (mup, mdn) upper and lower halves when the lower half matches
    list_chemical_symbols; False if not achieved within 5000 attempts.
"""
    fall=True
    contador=0
    while fall:
        mother2=mother1.copy()
        rv=randvector(deltadisp)
        mother2.translate(-rv)
        rotate_random(mother2)
        mup, mdn=get_gen(mother2)
        mdn_cs=mdn.get_chemical_symbols()
        mup_cs=mup.get_chemical_symbols()
        if (mdn_cs==list_chemical_symbols):
            fall=False
            return mup, mdn
        elif (mup_cs==list_chemical_symbols):
            fall=False
            positions = mup.get_positions()
            positions[:, 2] *= -1
            mup.set_positions(positions)
            positions = mdn.get_positions()
            positions[:, 2] *= -1
            mdn.set_positions(positions)
            return mdn, mup
        else:
            contador=contador+1
            if contador==5000:
                return False
#------------------------------------------------------------------------------------------
def mol_min_radius(atoms):
    """Compute 60% of the minimum distance from any atom to the center of mass.

in:
    atoms (ase.Atoms): cluster structure.
out:
    float: 0.6 times the distance from the center of mass to the nearest atom.
"""
    moleculeout=atoms.copy()
    moleculeout.translate(-moleculeout.get_center_of_mass())
    positions = atoms.get_positions()
    distances = np.linalg.norm(positions, axis=1)
    distances.sort()
    rmin=distances[0]*float(0.6)
    return rmin
#------------------------------------------------------------------------------------------
def complement(tot, part):
    """Return the multiset difference tot - part.

in:
    tot (list): full list of elements (e.g. chemical symbols).
    part (list): subset to remove.
out:
    list: copy of tot with one occurrence of each element in part removed.
"""
    ans=tot.copy()
    for elemi in part:
      if elemi in ans: ans.remove(elemi)
    return(ans)
#------------------------------------------------------------------------------------------
def cut_deavenho(singlefather, singlemother, atomlist):
    """Apply the Deaven-Ho cut-and-splice crossover to produce child fragments.

in:
    singlefather (ase.Atoms): first parent cluster.
    singlemother (ase.Atoms): second parent cluster.
    atomlist (list of str): full list of chemical symbols required in the child.
out:
    list of ase.Atoms: [upper_fragment_from_father, lower_fragment_from_mother]
    if a valid cut is found; False if it fails after trying half, half-1, and half+1 splits.
"""
    father1=singlefather.copy()
    mother1=singlemother.copy()
    father1.translate(-father1.get_center_of_mass())
    mother1.translate(-mother1.get_center_of_mass())
    rminf=mol_min_radius(father1)
    rminm=mol_min_radius(mother1)
    n0=len(atomlist)//2
    for ii, nff in enumerate([n0, n0-1, n0+1]):
        ans1=up_part_from_n(father1, nff, rminf)
        if (ans1 is False):
            return False
        fup, fdn  = ans1
        acfzlist=fup.get_chemical_symbols()
        rsdzlist=complement(atomlist, acfzlist)
        ans2=dn_part_from_zlist(mother1, rsdzlist, rminm)
        if (ans2 is False):
            if (ii == 2):
                return False
            else:
                continue
        else:
            mup, mdn = ans2
            break
    moleculeout=[fup, mdn]
    return moleculeout
#------------------------------------------------------------------------------------------
def crossover_deavenho(singlefather, singlemother, atomlist):
    """Perform a full Deaven-Ho crossover between two parent clusters.

in:
    singlefather (ase.Atoms): first parent cluster.
    singlemother (ase.Atoms): second parent cluster.
    atomlist (list of str): full list of chemical symbols required in the offspring.
out:
    ase.Atoms: child cluster formed by combining upper and lower halves from parents,
    rotated if needed to avoid atomic overlaps; False if all cut attempts fail.
"""
    moleculelist=Atoms()
    father=singlefather.copy()
    mother=singlemother.copy()
    result=cut_deavenho(singlefather, singlemother, atomlist)
    if not result:
        result=cut_deavenho(singlemother, singlefather, atomlist)
    if not result:
        return False
    else:
        gena,genb=result
        ratio_radii_crit=0.90
        var=gena.copy()
        moleculelist=[var+genb]
        dmin=distance_atoms2atoms(var, genb)
        test = False if (dmin < ratio_radii_crit) else True
        angle=0
        zeta=[0.0,0.0,1.0]
        vz=np.array([float(0.0), float(0.0), float(0.05)])
        while test == False:
            angle=angle+10
            var=rotate_vector_angle_deg(gena, zeta, angle)
            moleculelist.append(var+genb)
            dmin=distance_atoms2atoms(var, genb)
            if (dmin >= ratio_radii_crit):
                #writexyzs(moleculelist,'fenomeno.xyz')
                test=True
            if (angle >=360):
                angle=0
                gena.translate(-vz)
        moleculeout=moleculelist[-1]
        return moleculeout
#------------------------------------------------------------------------------------------
#def make_children(papy, mamy, atomlist):
#    moleculeout=[]
#    total_molecules=len(papy)
#    for ii in range(total_molecules):
#        ichildren=crossover_deavenho(papy[ii], mamy[ii], atomlist)
#        if ichildren is not False:
#            ichildren.e=float(0.0)
#            moleculeout.extend([ichildren])
#        else:
#            print("make_children(%s, %s) FALL" %(papy[ii].i,mamy[ii].i))
#    moleculeout=rename(moleculeout, 'pre_child_', 4)
#    return moleculeout
#------------------------------------------------------------------------------------------
def popgen_cukoo_childs(moleculein, memory_pre, memory_opt,index,number_of_childs,tol_gen,atomlist):
    """Generate a population of crossover offspring for one evolutionary generation.

in:
    moleculein (list of ase.Atoms): current population of optimized structures.
    memory_pre (list of ase.Atoms): archive of previously seen pre-optimization structures.
    memory_opt (list of ase.Atoms): archive of previously seen optimized structures.
    index (int): generation index used for labeling offspring.
    number_of_childs (int): desired number of crossover offspring.
    tol_gen (float): USR similarity tolerance for filtering duplicates.
    atomlist (list of str): chemical symbol list defining the required stoichiometry.
out:
    list of ase.Atoms: up to number_of_childs unique crossover offspring, renamed
    with 'mating_XX' prefix. Prints selection and count information.
"""
    if number_of_childs==0 or len(moleculein)<2:
        moleculeout=[]
        return moleculeout
    moleculeout=[]
    print("------------------------------------------------------------------")
    print("------------------------ MATING: CROSSOVER -----------------------")
    print("------------------------------------------------------------------")
    print("CONSTRUCTION OF THE PARENT LIST THROUGH ROULETTE")
    list_m=get_roulette_wheel_selection(moleculein, number_of_childs*2)
    list_p=get_roulette_wheel_selection(moleculein, number_of_childs*2)
    for i in range(number_of_childs*2):
        cross = crossover_deavenho(list_p[i], list_m[i], atomlist)
        cross.info['parents']={'p': list_p[i].info['i'], 'm':list_m[i].info['i']}
        cross.info['f_parents']={'p': list_p[i].info['f'], 'm': list_m[i].info['f']}
        if cross:  moleculeout.extend([cross])
    for atoms in moleculeout: atoms.info['e'] = 0
    moleculeout = comparator_usr_serial( moleculeout, tol_gen, 1)
    moleculeout = molin_sim_molref( moleculeout, memory_pre, tol_gen, 1)
    moleculeout = molin_sim_molref( moleculeout, memory_opt, tol_gen, 1)
    moleculeout=moleculeout[0:number_of_childs]
    moleculeout=rename( moleculeout,'mating_'+str(index).zfill(2),4)
    print("We have %d molecules type MATING from %d solicited" %(len(moleculeout), number_of_childs))
    if len(moleculeout)!=0:
        for imol in moleculeout:
            print("%s mating: %15s (f=%3.2f) and %15s (f=%3.2f)" %(imol.info['i'], imol.info['parents']['m'], imol.info['f_parents']['m'], imol.info['parents']['p'], imol.info['f_parents']['p']))
    return(moleculeout)
    # ordenar por energía el sumarry generatirion  ...
