import matplotlib.pyplot as plt
from leap_ec.probe import CartesianPhenotypePlotProbe
from leap_ec.representation import Representation

from leap_ec.individual import Individual
from leap_ec.algorithm import generational_ea

from leap_ec import ops
from leap_ec.decoder import IdentityDecoder
from leap_ec.real_rep.problems import CosineFamilyProblem
from leap_ec.real_rep.initializers import create_real_vector
from leap_ec.real_rep.ops import mutate_gaussian

# The fitness landscape
problem = CosineFamilyProblem(alpha=1.0, global_optima_counts=[2, 2], local_optima_counts=[2, 2])

# If no axis is provided, a new figure will be created for the probe to write to
trajectory_probe = CartesianPhenotypePlotProbe(contours=problem,
                                                xlim=(0, 1), ylim=(0, 1),
                                                granularity=0.025)

# Create an algorithm that contains the probe in the operator pipeline

pop_size = 100
ea = generational_ea(max_generations=20, pop_size=pop_size,
                     problem=problem,

                     representation=Representation(
                        individual_cls=Individual,
                        initialize=create_real_vector(bounds=[[0.4, 0.6]] * 2),
                        decoder=IdentityDecoder()
                     ),

                     pipeline=[
                         trajectory_probe,  # Insert the probe into the pipeline like so
                         ops.tournament_selection,
                         ops.clone,
                         mutate_gaussian(std=0.05, expected_num_mutations='isotropic', bounds=(0, 1)),
                         ops.evaluate,
                         ops.pool(size=pop_size)
                     ])
result = list(ea);