import emoo
import numpy as np
import pylab as pl

# Define the variables and their lower and upper search bounds
variables = [["x", -10, 10],
             ["y", -10, 10]]

# Define the list of objectives
objectives = ["f1", "f2"]

# This is the function which is going to be minimized
def func_to_optimize(parameters):
    
    x = parameters['x']
    y = parameters['y']
    
    f1 = 0.125*(20 + x + y)**2
    f2 = 100 + x*y

    return dict({"f1": f1, "f2": f2})

# After each generation this function is called
def checkpopulation(population, columns, gen):
   
    if gen == 0:
        i0 = np.where(population[:, columns['emoo_rank']] == 0)
        i1 = np.where(population[:, columns['emoo_rank']] == 1)
        i2 = np.where(population[:, columns['emoo_rank']] == 2)
        
        pl.plot(population[i0][:,columns['f1']], population[i0][:,columns['f2']], 'ro', label='pareto front 0')
        pl.plot(population[i1][:,columns['f1']], population[i1][:,columns['f2']], 'go', label='pareto front 1')
        pl.plot(population[i2][:,columns['f1']], population[i2][:,columns['f2']], 'bo', label='pareto front 2')
        
    if gen == 19:
        # finally we have to do a weigthing to select a single individual from the last generation
        dist0 = population[:,columns['f1']]**2 + population[:,columns['f2']]**2
        imin = np.argmin(dist0)
        
        best_individual = population[imin]
        x = best_individual[columns['x']]
        y = best_individual[columns['y']]
        f1 = best_individual[columns['f1']]
        f2 = best_individual[columns['f2']]
        
        print "x: %f; y: %.1f; f1: %.1f; f2: %.1f"%(x, y, f1, f2)
        
        pl.plot(population[:,columns['f1']], population[:,columns['f2']], 'x', label='final population')
        pl.xlabel("f1")
        pl.ylabel("f2")
        pl.legend()
        
        pl.savefig("example2.png")
        pl.show()
    
    print "Generation %d"%gen
        
# Initiate the Evlutionary Multiobjective Optimization
emoo = emoo.Emoo(N = 50, C = 100, variables = variables, objectives = objectives)
# Parameters:
# N: size of population
# C: size of capacity 

emoo.setup(eta_m_0 = 20, eta_c_0 = 20, p_m = 0.5)
# Parameters:
# eta_m_0, eta_c_0: defines the initial strength of the mution and crossover parameter (large values mean weak effect)
# p_m: probabily of mutation of a parameter (holds for each parameter independently)

emoo.get_objectives_error = func_to_optimize
emoo.checkpopulation = checkpopulation

emoo.evolution(generations = 20)