import functions
import numpy as np
import solvers

from matplotlib import pyplot as plt


plot_rest = True
f = functions.Rastrigin()
# f = functions.Sphere()
s = solvers.SimpleEvolutionStrategy(f, np.array([2, 2]))
old_fitness = 100
fitness = old_fitness * 0.9
old = None
plt.plasma()

while abs(old_fitness - fitness) > 0.001:
    old_fitness = fitness
    samples = s.sample(100)
    elite, fitness = s.rank(samples, 10)
    s.update(elite)

    if plot_rest:
        rest = np.setdiff1d(samples, elite, assume_unique=True)
        rest = rest.reshape((int(rest.shape[0]/2), 2))
    plt.pcolormesh(*f.grid())
    if plot_rest:
        plt.scatter(*rest.transpose(), color="dimgrey")
    if old is not None:
        plt.scatter(*old.transpose(), color="lightgrey")
    plt.scatter(*elite.transpose(), color="yellow")
    plt.scatter(*elite[0].transpose(), color="green")
    plt.show()
    old = elite