import random
import numpy as np
from ase.data              import atomic_numbers , covalent_radii 
from aegon.libcrossover    import get_gen
from aegon.libutils        import rename, writexyzs, rotate_vector_angle_deg
from aegon.libsel_roulette import get_roulette_wheel_selection
from aegon.libdiscusr      import molin_sim_molref, comparator_usr_serial
#-------------------------------------------------------------------------------
def separa_solapamientos(mvatoms,atoms, min_dist_factor=1.1, max_iter=50,doc='traj'):
    atoms_list_out=[]
    #atoms_list_out.extend([atoms])
    for _ in range(max_iter):
        moved = False
        positions = atoms.get_positions()
        num_atoms = len(atoms)
        for i in mvatoms:
            r1 = covalent_radii[atomic_numbers[atoms[i].symbol]]
            for j in range(num_atoms):
                if j != i:
                  dist_vec = positions[j] - positions[i]
                  r2 = covalent_radii[atomic_numbers[atoms[j].symbol]]
                  dist1 = np.linalg.norm(dist_vec)
                  dist2 = dist1/(r1 + r2)
                  if dist2 < min_dist_factor:
                      direction = dist_vec / dist1
                      move_vec = 0.5 * (min_dist_factor - dist2) * direction
                      positions[i] -= move_vec
                      if j in mvatoms:
                          positions[j] += move_vec
                      moved = True
        atoms.set_positions(positions)
        trial=atoms.copy()                #
        atoms_list_out.extend([trial])    #
        if moved is False:
            break
    writexyzs(atoms_list_out, doc) #
    return(atoms)
#-------------------------------------------------------------------------------
#MUTATION 1: Atoms displacement
def atom_displacement(moleculein, delta=0.4, min_dist_factor=1.1, max_iter=50):
    moleculeout = moleculein.copy()
    natoms = len(moleculeout)
    atoms_index = list(range(natoms))
    random.shuffle(atoms_index)
    num_mut = random.randint(1, natoms)
    mvatoms = atoms_index[:num_mut]
    #print(mvatoms)
    for ii in mvatoms:
        desplazamiento = np.random.uniform(-delta, delta, size=3)
        moleculeout.positions[ii] += desplazamiento
    moleculeout = separa_solapamientos(mvatoms,moleculeout, min_dist_factor, max_iter,doc='Mutation_one.xyz')
    moleculeout.info['o']={'i': moleculein.info['i'], 'f':moleculein.info.get('f',0)}
    return moleculeout
#----------------------------------------------------------------------
#MUTATION 2: twist mutation
def twist_mutation(atoms, delta=0.4, min_dist_factor=1.1, max_iter=50):
    atoms_mut = atoms.copy()
    atoms_mut.translate(-atoms_mut.get_center_of_mass())
    up, dn = get_gen(atoms_mut)
    up_indices=up.info['c']
    angulo = random.randint(5, 355)
    eje = [0, 0, 1]
    lista=[]
    for q in np.linspace(0, angulo, 50):
        upcopy=up.copy()
        up_rot=rotate_vector_angle_deg(upcopy, eje, q)
        atoms_mut=up_rot+dn
        atoms_mut.info['e']=atoms.info['e']
        atoms_mut.info['i']=atoms.info['i']
        lista.append(atoms_mut)
    writexyzs(lista, 'Mutation_two1.xyz') #
    atoms_mut = separa_solapamientos(up_indices, atoms_mut, min_dist_factor, max_iter, doc='Mutation_two2.xyz')
    atoms_mut.info['o']={'i': atoms.info['i'], 'f':atoms.info.get('f',0)}
    return atoms_mut
#------------------------------------------------------------------------------------------
def focused_expansion(molin, index, option='expansion'):
    """
    Vectorized version
    """
    molout = molin.copy()
    positions = molout.positions
    vector = positions[index].copy()
    vectors = positions - vector
    distances = np.linalg.norm(vectors, axis=1, keepdims=True)
    distances = np.where(distances == 0, 1e-10, distances)
    rnorm = vectors / distances
    signo=+1 if option == 'expansion' else -1
    edel = signo*np.exp(-distances) * 1.2
    new_positions = positions + rnorm * edel
    new_positions[index] = positions[index]
    molout.positions = new_positions
    # Genera los traj paara el 2
    return molout
#-------------------------------------------------------------------------
#MUTATION 4: atom permutation
def atom_permutation(moleculein):
    moleculeout = moleculein.copy()
    natoms = len(moleculeout)
    species = {}
    for i in range(natoms):species.setdefault(moleculeout.symbols[i], []).append(i)
    el_a, el_b = random.sample(list(species.keys()), 2)
    grupo_a = species[el_a]
    grupo_b = species[el_b]
    ia = random.choice(grupo_a)
    ib = random.choice(grupo_b)
    moleculeout.symbols[ia], moleculeout.symbols[ib] = moleculeout.symbols[ib], moleculeout.symbols[ia]
    ra = covalent_radii[atomic_numbers[moleculeout[ia].symbol]]
    rb = covalent_radii[atomic_numbers[moleculeout[ib].symbol]]
    opt_a, opt_b = ('expansion', 'impasion') if ra > rb else ('impasion', 'expansion')
    moleculeout = focused_expansion(moleculeout, ia, option=opt_a)
    moleculeout = focused_expansion(moleculeout, ib, option=opt_b)
    moleculeout.info['o']={'i': moleculein.info['i'], 'f':moleculein.info.get('f',0)}
    return moleculeout
#--------------------------------------------------------------------------
def inequivalent_different_finder(moleculein, tolerance=0.3, eigen_tolerance=0.008, matrix_tolerance=0.1):
    moleculeout = moleculein.copy()
    natoms = len(moleculeout)
    list_ineq   = inequivalent_finder(moleculeout, tolerance, eigen_tolerance, matrix_tolerance)
    if len(list_ineq) == 0:
        return(False)
    else:
        species = {}
        for i in list_ineq:species.setdefault(moleculeout.symbols[i], []).append(i)
        return(species)
#-------------------------------------------------------------------------------
def make_mutants(the_chosen_ones_list):
    moleculeout=[]
    for imol in the_chosen_ones_list:
        especies= list(set(imol.symbols))
        wich = random.choice([0,1,2]) if len(especies)>1 else random.choice([0,1])
        if wich == 0:
            muty = atom_displacement(imol, delta=0.4, min_dist_factor=1.1, max_iter=50)
            muty.info['Mutation'] = 'Atom Displacement'
        if wich == 1:
            muty = twist_mutation(imol, delta=0.4, min_dist_factor=1.1, max_iter=50)
            muty.info['Mutation'] = 'Twist Mutation'
        if wich == 2:
            muty = atom_permutation(imol)
            muty.info['Mutation'] = 'Atom Permutation'
        moleculeout.append(muty)
    moleculeout = rename(moleculeout,'pre_mutant_',4)
    return moleculeout
#-------------------------------------------------------------------------------
def popgen_cukoo_mutis(moleculein,memory_pre, memory_opt,index,number_of_mutants,tol_gen):
    if number_of_mutants == 0:
        moleculeout = []
        return moleculeout
    #silence = 1 if verbose == 0 else 0 # ¿Que es verbose?
    print("------------------------------------------------------------------")
    print("-------------------------GENERATOR OF MUTANTS---------------------")
    print("------------------------------------------------------------------")
    the_chosen_ones = get_roulette_wheel_selection(moleculein, number_of_mutants * 10)
    moleculeout = make_mutants(the_chosen_ones)
    #print ("Compare with itself: %d" % (len(moleculeout)))
    moleculeout = comparator_usr_serial(moleculeout, tol_gen, 1)
    #print ("Compare with memory_pre: %d" % (len(moleculeout)))
    moleculeout = molin_sim_molref(moleculeout, memory_pre, tol_gen, 1)
    #print ("Compare with memory_opt: %d" % (len(moleculeout)))
    moleculeout = molin_sim_molref(moleculeout, memory_opt, tol_gen, 1)
    moleculeout = moleculeout[0:number_of_mutants]
    moleculeout = rename(moleculeout, 'mutant_'+str(index).zfill(2),4)
    print ("We have %d molecules type MUTANT from %d solicited." % (len(moleculeout), number_of_mutants))
    if len(moleculeout)!=0:
        for imol in moleculeout:
            print("%s by: %-18s in %15s (f=%3.2f)" % (imol.info['i'],imol.info['Mutation'],imol.info['o']['i'],imol.info['o']['f']))
    return moleculeout
#-----------------------------------------------------------------------
def popgen_cukoo_childs(moleculein, memory_pre, memory_opt,index,number_of_childs,tol_gen,atomlist):
    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  ...