import matplotlib.pyplot as plt
from leap_ec.probe import SumPhenotypePlotProbe
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.binary_rep.problems import DeceptiveTrap
from leap_ec.binary_rep.initializers import create_binary_sequence
from leap_ec.binary_rep.ops import mutate_bitflip

# The fitness landscape
problem = DeceptiveTrap()

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

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

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

                     representation=Representation(
                        individual_cls=Individual,
                        initialize=create_binary_sequence(length=dimensions)
                     ),

                     pipeline=[
                         trajectory_probe,  # Insert the probe into the pipeline like so
                         ops.tournament_selection,
                         ops.clone,
                         mutate_bitflip(expected_num_mutations=1),
                         ops.evaluate,
                         ops.pool(size=pop_size)
                     ])
result = list(ea);