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, rename
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'):
    """Iteratively separate overlapping atoms by pushing them apart.

in:
    mvatoms (list of int): indices of atoms that are allowed to move.
    atoms (ase.Atoms): structure containing overlapping atoms (modified in-place).
    min_dist_factor (float): minimum allowed interatomic distance as a fraction
        of the sum of covalent radii (default 1.1).
    max_iter (int): maximum number of separation iterations (default 50).
    doc (str): filename for the trajectory XYZ output written at each iteration.
out:
    ase.Atoms: structure with overlaps resolved. Each iteration is written to doc.
"""
    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):
    """Apply mutation 1: randomly displace a random subset of atoms.

in:
    moleculein (ase.Atoms): structure to mutate. Must have info['i'].
    delta (float): maximum displacement magnitude in Angstroms (default 0.4).
    min_dist_factor (float): overlap resolution threshold (default 1.1).
    max_iter (int): maximum overlap resolution iterations (default 50).
out:
    ase.Atoms: mutated structure with info['o'] recording parent name and fitness.
    Overlaps are resolved after displacement. Trajectory written to 'Mutation_one.xyz'.
"""
    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):
    """Apply mutation 2: rotate the upper half of a cluster around the Z axis.

in:
    atoms (ase.Atoms): structure to mutate. Must have info['e'] and info['i'].
    delta (float): unused displacement parameter (kept for API consistency).
    min_dist_factor (float): overlap resolution threshold (default 1.1).
    max_iter (int): maximum overlap resolution iterations (default 50).
out:
    ase.Atoms: mutated structure with a random twist angle applied to the upper
    half. Trajectory written to 'Mutation_two1.xyz' and 'Mutation_two2.xyz'.
    info['o'] records parent name and fitness.
"""
    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'):
    """Apply a focused radial expansion or contraction around a given atom.

in:
    molin (ase.Atoms): structure to modify.
    index (int): index of the pivot atom; this atom remains fixed.
    option (str): 'expansion' to push atoms outward, 'impasion' to pull inward.
out:
    ase.Atoms: structure with all other atoms displaced radially from the pivot
    by an exponentially decaying amount.
"""
    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):
    """Apply mutation 4: swap the identities of two atoms from different species.

in:
    moleculein (ase.Atoms): multi-element structure to mutate. Must have at least
        two distinct element types and info['i'].
out:
    ase.Atoms: mutated structure with two atoms swapped, followed by focused
    expansion/contraction to accommodate the size difference. info['o'] records
    parent name and fitness.
"""
    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):
    """Find symmetry-inequivalent atoms grouped by species for targeted mutation.

in:
    moleculein (ase.Atoms): structure to analyze.
    tolerance (float): distance tolerance 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:
    dict or False: dictionary mapping element symbols to lists of inequivalent atom
    indices, or False if no inequivalent atoms are found.
"""
    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):
    """Apply a random mutation operator to each structure in the input list.

in:
    the_chosen_ones_list (list of ase.Atoms): structures selected for mutation.
        Each must have info['i'] and optionally info['f'].
out:
    list of ase.Atoms: mutated structures renamed with 'pre_mutant_' prefix.
    Each has info['Mutation'] recording the applied operator name.
    Operators: 'Atom Displacement', 'Twist Mutation', 'Atom Permutation'.
"""
    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):
    """Generate a population of mutant structures 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 naming the mutants.
    number_of_mutants (int): desired number of unique mutant structures.
    tol_gen (float): USR similarity tolerance for filtering duplicate mutants.
out:
    list of ase.Atoms: up to number_of_mutants unique mutants, renamed with
    'mutant_XX' prefix. Prints selection and count information.
"""
    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
