import functions
import numpy as np
import solvers

from matplotlib import pyplot as plt


def main():
    plot_rest = True
    f = functions.Rastrigin()
    # f = functions.Sphere()
    # s = solvers.SimpleEvolutionStrategy(f, np.array([2, 2]))
    s = solvers.CMAEvolutionStrategy(f, np.array([2.0, 2.0]), np.array([[2.0, 0.0], [0.0, 2.0]]))
    old_fitness = 100
    fitness = old_fitness * 0.9
    old = None
    plt.set_cmap('Spectral')

    while abs(old_fitness - fitness) > 0.001:
        old_fitness = fitness
        samples = s.sample(2000)
        elite, fitness = s.rank(samples, 150)
        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="orange")
        plt.scatter(*elite.transpose(), color="yellow")
        # plt.scatter(*elite[0].transpose(), color="green")
        print('old fitness: {}\nnew fitness: {}\nimprovement: {}\n'.format(old_fitness, fitness,
            old_fitness - fitness))
        plt.show()
        old = elite 


if __name__ == '__main__':
    main()