As library

####
#### when you change set_logger(), you need to restart the kernel
####
from cryspy.util.utility import set_logger    # optional
set_logger()    # optional
#set_logger(noprint=True, logfile='log_cryspy', errfile='err_cryspy')    # write log and err messages to files

from cryspy.RS.gen_struc_RS import gen_pyxtal

nstruc = 10
atype = ('Na', 'Cl')
nat = (4, 4)
mindist = ((2.0, 1.5),
           (1.5, 2.0))
spgnum = 'all'

init_struc_data = gen_pyxtal.gen_struc(
    nstruc=nstruc,
    atype=atype,
    nat=nat,
    mindist=mindist,
    spgnum=spgnum,
)

import pickle
with open('./Cu4Au4_sample/opt_struc_data.pkl', 'rb') as f:
    opt_struc_data = pickle.load(f)

from cryspy.EA.gen_struc_EA import crossover

# you can change parent_A and parent_B
parent_A = opt_struc_data[0]
parent_B = opt_struc_data[1]

atype = ('Cu', 'Au')
nat = (4, 4)
mindist = ((1.5, 1.5),
           (1.5, 1.5))

child = crossover.gen_child(
    atype=atype,
    nat=nat,
    mindist=mindist,
    parent_A=parent_A,
    parent_B=parent_B,
)

# child: pymatgen Structure

from cryspy.EA.gen_struc_EA import permutation

# you can change parent_A
parent_A = opt_struc_data[0]

atype = ('Cu', 'Au')
nat = (4, 4)
mindist = ((1.5, 1.5),
           (1.5, 1.5))
ntimes = 1    # number of times to perform permutatio

child = permutation.gen_child(
    atype=atype,
    mindist=mindist,
    parent_A=parent_A,
    ntimes=ntimes,
)

# child: pymatgen Structure

from cryspy.EA.gen_struc_EA import strain

atype = ('Cu', 'Au')
nat = (4, 4)
mindist = ((1.5, 1.5),
           (1.5, 1.5))
sigma_st = 0.05    # standard deviation of strain

child = strain.gen_child(
    atype=atype,
    mindist=mindist,
    parent_A=parent_A,
    sigma_st=sigma_st,
)

import pickle
with open('./Cu4Au4_sample/opt_struc_data.pkl', 'rb') as f:
    opt_struc_data = pickle.load(f)
with open('./Cu4Au4_sample/rslt_data.pkl', 'rb') as f:
    rslt_data = pickle.load(f)

struc_data = opt_struc_data    # dict
fitness = rslt_data['E_eV_atom'].to_dict()    # you may include None or np.nan for values

from cryspy.EA.survival import survival_fittest
from cryspy.EA.gen_struc_EA.select_parents import SelectParents
from cryspy.EA.gen_struc_EA import crossover, permutation, strain

n_fittest = 5    # number of survivors

ranking, _, _ = survival_fittest(
    fitness=fitness,
    struc_data=struc_data,
    elite_struc=None,
    elite_fitness=None,
    n_fittest=n_fittest,
    fit_reverse=False,
    emax_ea=None,
    emin_ea=None,
)

# ranking <-- e.g. [2, 1, 0, 7, 9] without structure duplicaiton

sp = SelectParents(ranking)    # after set_xxx, we can use sp.get_parents(n_parent)
sp.set_tournament(t_size=2)

atype = ('Cu', 'Au')
nat = (4, 4)
mindist = ((1.5, 1.5),
           (1.5, 1.5))
n_crsov = 5    # number of structures to be generated by crossover
#id_start = len(init_struc_data)  # next Structure ID
id_start = 10

co_children, co_parents, co_operation = crossover.gen_crossover(
    atype=atype,
    nat=nat,
    mindist=mindist,
    struc_data=struc_data,
    sp=sp,
    n_crsov=n_crsov,
    id_start=id_start,
)

# co_children <-- dict {ID: pymatgen Structure, ID: pymatgen Structure, ...}
# co_parents  <-- e.g. {10: (2, 7), 11: (2, 1), 12: (2, 1), 13: (0, 2), 14: (2, 1)}
# co_operation <-- e.g. {10: 'crossover', 11: 'crossover', ...}

n_perm = 5    # number of structures to be generated by permutation
#id_start = len(init_struc_data) + n_crsov   # next Structure ID
id_start = 15
ntimes = 1    # number of times to perform permutation

pm_children, pm_parents, pm_operation = permutation.gen_permutation(
    atype=atype,
    mindist=mindist,
    struc_data=struc_data,
    sp=sp,
    n_perm=n_perm,
    id_start=id_start,
    ntimes=ntimes,
)

# pm_children <-- dict {ID: pymatgen Structure, ID: pymatgen Structure, ...}
# pm_parents  <-- e.g. {15: (2,), 16: (1,), 17: (2,), 18: (1,), 19: (1,)}
# pm_operation <-- e.g. {15: 'permutaion', 16: 'permutation', ...}

n_strain = 5    # number of structures to be generated by strain
#id_start = len(init_struc_data) + n_crsov + n_perm   # next Structure ID
id_start = 20
sigma_st = 0.05    # standard deviation of strain

st_children, st_parents, st_operation = strain.gen_strain(
    atype=atype,
    mindist=mindist,
    struc_data=struc_data,
    sp=sp,
    n_strain=n_strain,
    id_start=id_start,
    sigma_st=sigma_st,
)

# st_children <-- dict {ID: pymatgen Structure, ID: pymatgen Structure, ...}
# st_parents  <-- e.g. {20: (1,), 21: (2,), 22: (0,), 23: (2,), 24: (2,)}
# st_operation <-- e.g. {20: 'strain', 21: 'strain', ...}