from ..options import Options
from .evolutionaryalgorithm import EvolutionaryAlgorithm
from .dna import dna
from ..tex import tex

class GA(EvolutionaryAlgorithm):
    _def_opts = {
        'dna_size': lambda val: isinstance(val, int) and val > 0,
        'population_size': lambda val: isinstance(val, int) and val > 0,
        'point_mutation_chance': (0.02, slice(0, 1)),
        'point_mutation_strength': (0.05, slice(0, 1)),
        'selection_std_dev': (0.30, slice(0, 1)),
        'keep_elites_rate': (0.05, slice(0, 10**10)),
        'mutate_elites': False
    }

    @staticmethod
    def default_options(dna_size=None, population_size=None):
        opts = Options(GA._def_opts)

        if dna_size is not None:
            opts.dna_size = dna_size
        if population_size is not None:
            opts.population_size = population_size

        return opts

    def __init__(self, fitness_function, options, seed_factory=None):
        super().__init__(
            options.dna_size,
            options.population_size,
            fitness_function,
            seed_factory)

        self.point_mutation_chance = options.point_mutation_chance
        self.point_mutation_strength = options.point_mutation_strength
        self.selection_std_dev = options.selection_std_dev
        self.keep_elites_rate = options.keep_elites_rate
        self.mutate_elites = bool(options.mutate_elites)

    def do_step(self):
        # crossover
        elites, non_elites = self.population.make_crossovered_dna(
            self.selection_std_dev,
            self.keep_elites_rate)

        # point mutation:
        if elites is not None and self.mutate_elites:
            non_elites = dna.mutate(
                tex.join(1, elites, non_elites),
                self.point_mutation_chance,
                self.point_mutation_strength)
            elites = None
        else:
            non_elites = dna.mutate(
                non_elites,
                self.point_mutation_chance,
                self.point_mutation_strength)

        # new population
        new_pop = self.create_population()
        new_pop.push((elites, non_elites))
        self.set_populaion(new_pop)