import numpy as np
from ase import Atoms
from aegon.libsel_roulette import get_roulette_wheel_selection
from aegon.libutils        import random_vector, rotate_random, distance_atoms2atoms, rotate_vector_angle_deg, rename
from aegon.libdisc_usr     import molin_sim_molref, comparator_usr_serial
#------------------------------------------------------------------------------------------
""" CUT IN THE Z AXIS """
def get_gen(atoms):
    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 up_part_from_n(father1, nff, deltadisp):
    fall=True
    contador=0
    while fall:
        father2=father1.copy()
        rv=random_vector(radius=deltadisp, random_radius=True)
        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):
    fall=True
    contador=0
    while fall:
        mother2=mother1.copy()
        rv=random_vector(radius=deltadisp, random_radius=True)
        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):
    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):
    ans=tot.copy()
    for elemi in part:
      if elemi in ans: ans.remove(elemi)
    return(ans)
#------------------------------------------------------------------------------------------
def cut_deavenho(singlefather, singlemother, atomlist):
    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):
    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 popgen_cukoo_childs(moleculein, memory_pre, memory_opt,index,number_of_childs,tolsim,tolener,mono,atomlist):
    if number_of_childs==0 or len(moleculein)<2: return []
    moleculeout=[]
    print("------------------------------------------------------------------")
    print("------------------------ MATING: CROSSOVER -----------------------")
    print("------------------------------------------------------------------")
    print("CONSTRUCTION OF THE PARENT LIST THROUGH ROULETTE")
    list_m=get_roulette_wheel_selection(moleculein, number_of_childs*5)
    list_p=get_roulette_wheel_selection(moleculein, number_of_childs*5)
    moleculeout,new_cross = [],[]
    for i in range(len(list_p)):
        cross_imol = crossover_deavenho(list_p[i], list_m[i], atomlist)
        cross_imol.info['parents']={'p': list_p[i].info['i'], 'm':list_m[i].info['i']}
        cross_imol.info['f_parents']={'p': list_p[i].info['f'], 'm': list_m[i].info['f']}
        cross_imol.info['e'] = 0
        if cross_imol:
            if len(new_cross)!= 0:  #Está bien
                original_cross = molin_sim_molref([cross_imol],new_cross, tolsim, tolener, mono, 0)
                if original_cross: new_cross.append(cross_imol)
                else: continue
            else: new_cross.append(cross_imol)
            cross_imol = molin_sim_molref([cross_imol], memory_pre, tolsim, tolener, mono, 0)
            cross_imol = molin_sim_molref(cross_imol, memory_opt, tolsim, tolener, mono, 0)
            if not cross_imol: continue
            moleculeout.extend(cross_imol)
            if len(moleculeout) == number_of_childs:break
    print("We have %d molecules type MATING from %d solicited" %(len(moleculeout), number_of_childs))
    if moleculeout:
      moleculeout=rename(moleculeout,'mating_'+str(index).zfill(2),4)
      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']))
    if not moleculeout: print('The maximum number of attempts to create new crosses has been exceeded.')
    return(moleculeout)
