from aegon.libutils        import sort_by_energy, cutter_energy, rename, writexyzs
from aegon.libmolgen       import make_molecules_random
from aegon.libcrossover    import popgen_cukoo_childs
from aegon.libmutants      import popgen_cukoo_mutis
from aegon.libstdio        import read_main_input
from aegon.libcalc_lj      import opt_LJ_parallel
from aegon.libsel_roulette import get_fitness
from aegon.libdisc_usr     import comparator_usr_serial, molin_sim_molref
tolsij=0.95
tolene=0.10
#------------------------------------------------------------------------------------------
def display_mol_info(moleculein, flagsum=0, ngen=0):
    if len(moleculein)==0:
        print("\n------------ALL MOLECULES DISCRIMINATED. GLOMOS FINISH.------------")
    if flagsum == 0:
        print("-----------------------GENERATION %d SUMMARY-----------------------" %(ngen))
    else:
        print("--------------------------GLOBAL SUMMARY---------------------------")
    fitness=get_fitness(moleculein)
    for atoms, f in zip(moleculein, fitness): atoms.info['f'] = f
    for ii, imol in enumerate(moleculein):
        deltae=imol.info['e'] - moleculein[0].info['e']
        jj=str(ii+1).zfill(5)
        if flagsum == 0:
            print("#%s %-14s with %13.8f eV (%10.8f)" %(jj, imol.info['i'], imol.info['e'], deltae))
        else:
            print("#%s %-14s with %-13.8f eV (%-10.8f) (f=%-3.2f)" %(jj, imol.info['i'], imol.info['e'], deltae, imol.info['f']))
    if flagsum != 0 : print('')
#------------------------------------------------------------------------------------------
def genetic_algorithm(inputfile):
    df = read_main_input(inputfile)
    composition = df.get_comp(key='COMPOSITION')
    atomlist=composition.atoms
    nameid=composition.name
    mono=True if len(composition.elements)==1 else False
    ninitpop=df.get_int(key='nof_initpop', default=10)
    nmatings=df.get_int(key='nof_matings', default=5)
    ngenerations=df.get_int(key='nof_generations', default=5)
    tolener=df.get_float(key='tol_energy', default=tolene)
    tolsim=df.get_float(key='tol_similarity', default=tolsij)
    ecut=df.get_float(key='cutoff_energy', default=10.0)
    npop=df.get_int(key='cutoff_population', default=10)
    nproc=df.get_int(key='nof_processes', default=2)
    nmut = df.get_int(key='nof_mutants',default=5)
    nisomers= df.get_int(key='nof_repeats',default=2)
    ncycle= df.get_int(key='nof_stagnant',default=3)
    print('-------- Global Optimization of Molecular Systems (GLOMOS) --------')
    #print('Input file           = %s' %(inputfile))
    print('Chemical Formula     = %s\n' %(nameid))
    print('EVOLUTIVE PARAMETERS:')
    #print('Option               = Modified Genetic Algorithm')
    print('Initial Population   = %d' %(ninitpop))
    print('Number of matings    = %d' %(nmatings))
    print('Number of mutants    = %d\n' %(nmut))
    print('DISCRIMINATION PARAMETERS:')
    print('Tol for energy       = %4.2f' %(tolener))
    print('Tol for similarity   = %4.2f' %(tolsim))
    print('Energy Cut-off       = %f' %(ecut))
    print('Max population size  = %d\n' %(npop))
    print('STOP CRITERION:')
    print('Max generations       = %d' %(ngenerations))
    print('Max repeated isomers  = %d' %(nisomers))
    print('Max stagnant cycles   = %d\n' %(ncycle))
    print('THEORY LEVEL: LJ, SC (Cu, Pd, Ag, Pt, Au, Pb, Al)')
    print('Number of parallel local opts= %d\n\n' %(nproc))
    xrand = make_molecules_random(atomlist, ninitpop)
    xrand=rename(xrand,'random_00',4)
    #xopt=[opt_EMT(ix) for ix in xrand]
    xopt=opt_LJ_parallel(xrand, nproc)
    xopt=cutter_energy(xopt, ecut)
    xopt_sort=sort_by_energy(xopt, 1)
    xopt_sort=comparator_usr_serial(xopt_sort, tolsim, mono)
    memory_opt = xopt_sort.copy()
    memory_pre = xrand.copy()
    xopt_sort=xopt_sort[:npop]
    display_mol_info(xopt_sort, flagsum=1, ngen=0)
    iter_max = 2
    prev_xopt = None
    print('')
    for igen in range(ngenerations):
        #-----------------------------------------------------------------------
        current_xopt = xopt_sort[:nisomers]
        if prev_xopt is not None and current_xopt == prev_xopt:
            counter += 1
            if counter == ncycle:
                print(f'The user requested to stop the program in {ngenerations} consecutive generations (crit_stop_nrep),')
                print(f'with the same global minimum. However, all optimized individuals in the first {nisomers} converge \n to similar individuals from the previous {ncycle} generations, so GLOMOS decided to STOP.')
                writexyzs(memory_opt,'memory.xyz')
                return xopt_sort
        else: counter = 1
        prev_xopt = current_xopt
        for ii in range(iter_max):
            print("---------------------------GENERATION %d--------------------------" %(igen+1))
            generation_pre,generation_opt = [],[]
            crosses_pre = popgen_cukoo_childs(xopt_sort, memory_pre, memory_opt,igen+1,nmatings,tolsim,tolener,mono,atomlist)
            if not crosses_pre and nmatings != 0: print('CROSS FILTER: TYPE 1: No non-optimized crosses were obtained in generation %d.' %(igen+1))
            mutant_pre = popgen_cukoo_mutis(xopt_sort, memory_pre, memory_opt, igen+1, nmut, tolsim,tolener,mono)
            if not mutant_pre and nmut != 0: print('MUTANT FILTER: TYPE 1: No new mutants were created in generation %d.' %(igen+1))
            generation_pre = crosses_pre + mutant_pre
            if generation_pre:
                print('')
                memory_pre.extend(generation_pre)
                generation_opt=opt_LJ_parallel(generation_pre, nproc)
                generation_opt=cutter_energy(generation_opt, ecut)
                generation_opt=sort_by_energy(generation_opt,1)
                generation_opt=comparator_usr_serial(generation_opt, tolsim, tolener, mono)
                generation_opt=molin_sim_molref(generation_opt, memory_opt, tolsim,tolener, mono)
                if generation_opt:
                    memory_opt.extend(generation_opt)
                if not generation_opt:
                    print('------------------------------------------------------------------\n')
                    if ii != iter_max-1: print(f'No individuals were obtained in generation {igen+1}. The process is repeated.\n')
                    continue
            if not generation_pre:
                print('------------------------------------------------------------------\n')
                if ii != iter_max-1: print(f'No individuals were obtained in generation {igen+1}. The process is repeated.\n')
                continue
            generation_opt=sort_by_energy(generation_opt,1)
            display_mol_info(generation_opt, flagsum=0, ngen=igen+1)
            xopt_sort=sort_by_energy(xopt_sort+generation_opt, 1)
            xopt_sort=xopt_sort[:npop]
            display_mol_info(xopt_sort, flagsum=1)
            break
        else:
            print(f'CAUTION!: Gen {igen+1} | Attempts exhausted. The procedure is terminated.\n')
            display_mol_info(xopt_sort, flagsum=1)
            break
    writexyzs(memory_opt,'memory.xyz')
    return xopt_sort
