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

f = plt.figure()  # Setup a figure to plot to
plot_probe = FitnessPlotProbe(ylim=(0, 70), ax=plt.gca())


# Create an algorithm that contains the probe in the operator pipeline
from leap_ec.individual import Individual
from leap_ec.decoder import IdentityDecoder
from leap_ec import ops
from leap_ec.real_rep.problems import SpheroidProblem
from leap_ec.real_rep.ops import mutate_gaussian
from leap_ec.real_rep.initializers import create_real_vector

from leap_ec.algorithm import generational_ea

l = 10
pop_size = 10
ea = generational_ea(max_generations=100, pop_size=pop_size,
                     problem=SpheroidProblem(maximize=False),

                     representation=Representation(
                        individual_cls=Individual,
                        decoder=IdentityDecoder(),
                        initialize=create_real_vector(bounds=[[-5.12, 5.12]] * l)
                     ),

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