# %matplotlib notebook
import numpy as np
import matplotlib.pyplot as plt
import sys
from mpl_toolkits.mplot3d import Axes3D
import math

#MOO Problema de Optimización Multi Objetivo
#Primero se implementa la clase MOO_Problema, que permite cualquier número de funciones. 
#Cada función es almacenada en una matriz (Nxd) de numpy, donde (d) es la dimensión de la variable de decisión.

class MOO_Problem():
    ''' 
    Clase MOO_Problem:  define las diferentes funciones de restricción que se necesitan minimizar. 
    Parametros utilizados: 
        d: Dimensión del vector de entrada para las funciones 
        inferior: (mx1) matriz numpy que proporciona el límite inferior de las m funciones. Predeterminado = -1 
        superior: (mx1) matriz numpy que proporciona el límite superior de las m funciones. Predeterminado = 1 
        *args: numero m de funciones python capaces de ejecutarse en (Nxd) matrices numpy MOO_Problem(func1, func2, func3, ...) 
    ''' 

    #Constructor de la clase MOO_Problem
    def __init__(self, *args, d=None,lower=None, upper=None):
        if d is None:
            print("Falta el parámetro obligatorio 'd'")
            sys.exit(-1)
        
        #Se asigna variables que entran a variables globales en la clase
        self.d = d
        self.m = len(args)

        #Validando que exista al menos una funcion
        if self.m < 1:
            print("Ingresa al menos una función objetivo")
            sys.exit(-1)
        
        #Se crea un arreglo de funciones
        self.functions = list(args)

        #Se crea una matriz de dimension [1,d] con los valores que se reciben de lower y upper, en caso contrario se llena de -1's si es lower y 1's si es upper
        self.lower = lower if lower is not None else -np.ones((1, self.d))
        self.upper = upper if upper is not None else np.ones((1, self.d))

    def evaluate(self, x):
        ''' 
        Evalua las funciones en múltiples vectores posibles 
        : x: (N, m) matriz numpy donde N es el número de vectores y m es el número de funciones
        Devuelve: (N, m) matriz numpy de resultados evaluados 
        '''
        #N =  Tamaño de la poblacion
        #m = Numero de funciones objetivo
        N = x.shape[0] 
        obj = np.empty((N, self.m)) #Crea matriz vacia de [2,10]
        for i in range(self.m):
            aux = np.squeeze(self.functions[i](x)) #Evalua todos los individuos de la poblacion por cada función objetivo y genera un vector por funcion con squeeze 
            obj[:, i] = aux #Crea una matriz donde cada columa son los valores de la evaluacion de cada funcion objetivo

        #print("objeto3=",obj)
        return obj
'''
A continuación, se implementa la clase NSGAII. 
Los parámetros utilizados en esta clase son:

* moo_problem: El problema MOO que necesita ser optimizado. Se utiliza el supuesto de que todas las funciones deben minimizarse
* pop_size: Número de individuos en la población por cada iteración
* p_mutate: Probabilidad de mutación
* p_crossover: Probabilidad de cruce
* n_iter: Número máximo de iteraciones
* max_evals: Número máximo de evaluaciones para las funciones
* initial_pop: Población inicial para comenzar el algoritmo. Si no se proporciona, se genera al azar.

p_mutate/ p_crossover puede afectar significativamente al algoritmo. A continuación, se grafica el cambio en la descendencia con respecto a sus padres debido al cruce, 
para valores variables de p_crossover.

Effect of p_crossover
	Otro conjunto de parámetros que no están disponibles a través de la clase, pero que se pueden editar directamente dentro de la función NSGAII.evolve()es dis_c/ dis_m. 
	A continuación se muestra la distribución del mismo cambio en función de diferentes valores de dis_ccon p_crossover = 1.

Effect of dis_c
	Para todo el algoritmo, observamos la tendencia de que para valores muy altos de dis_c/ dis_m, el óptimo de Pareto logrado es menos diverso.
'''

class NSGAII():
    '''
    Clase para ejecutar el algoritmo NSGA II 
    '''
    def __init__(self, moo_problem, pop_size=100 , p_mutate=1.0, p_crossover=1.0, n_iter=100, max_evals= 100*500, initial_pop=None):
        ''' 
        Param: 
            moo_problem: Problema multiobjetivo de la clase MOO_Problem
            pop_size: Número de individuos de la población por cada iteración 
            p_mutate: probabilidad de mutación (entre 0 y 1, por defecto 1) 
            p_crossover: probabilidad de cruce (entre 0 y 1, por defecto 1) 
            n_iter: Número máximo de iteraciones permitidas (por defecto = 100) 
            max_evals: número máximo de evaluaciones de funciones permitidas (predeterminado = 500 por individuo) 
            initial_pop: (Nxd) matriz numérica de la población inicial de individuos. Si no hay ninguno, se inicializó aleatoriamente 
        ''' 

        #Se asignan variables que entran a variables globales en la clase
        self.N = pop_size #Tamaño de la población
        self.p_m = p_mutate #Porcentaje de mutación
        self.p_c = p_crossover #Porcentaje de cruza
        self.max_iter = n_iter #Maximo número de iteraciones
        self.max_evals = max_evals
        self.moo = moo_problem

        if initial_pop is not None and initial_pop.shape[1]==self.moo.d:
            self.pop = [initial_pop, self.moo.evaluate(initial_pop)]
            print("Población inicializada")
            #input(".....")
        else:
            #En caso de que no se proporcione una poblacion inicial se genera una poblacion  con valores  aleatorios
            #if initial_pop is not None:
                #print("La población inicial no contiene elementos. Inicializando población aleatoria...")
                #input(".....")
            
            print("La población inicial no contiene elementos. Inicializando población aleatoria...")

            #Se crea un arreglo de valores aleatorios de tamaño [N,m] y se multiplica por (1-1) + (-1)
            random_pop = np.random.random((self.N, self.moo.d)) * (self.moo.upper - self.moo.lower) + self.moo.lower
            
            #Evaluando a la poblacion incial (random_pop) y crea una matriz donde la primer columa son los valores de andom_pop y el resto de columnas son el resultado de evaluar cada funcion
            self.pop = [random_pop, self.moo.evaluate(random_pop)]

    def non_dominated_sort(self, n_sort=None):
        '''
        Realiza una clasificación no dominada de la población en esta iteración 
        Parámetros: 
        	n_clasificación: número máximo de individuos a clasificar 
        
        Devolución: 
            front_ids: (Nx1) vector donde front_id[i] es el número principal de pop[i] 
            max_front: Número de frentes calculados
        '''
        # Inicialización
        #print("costo pop--",self.pop[1])

        pop_cost = self.pop[1]  #Matriz de funciones objetivo evaludas, cada columna son los valores de una función objetivo

        N = pop_cost.shape[0] #Numero de  individuos en la poblacion

        #La función np.unique() recupera todos los indices de los valores únicos en la columna pop_cost[:,0] (primera columna de la funcion objetivo 1)
        _, loc = np.unique(pop_cost[:,0], return_inverse=True)


        #La primera vez que se ejecuta "n_sort = None"
        if n_sort is None: n_sort = len(loc) #n_sort toma el valor del número total de individuos en (loc)
        sorted_cost = pop_cost[pop_cost[:,0].argsort(), :]  #Ordena los valores de columna de cada funcion objetivo de mayor a menor

        #print("pop cost:",pop_cost)
        #print("sorted_cost:",sorted_cost)

        front_id = np.inf*np.ones(N) #Crea un vector con N elementos llamados "inf"

        #print(front_id)

        max_front = 0 #Establece el valor inicial máximo de frontera

        #Ordenamiento no dominado
        while np.sum(front_id < np.inf) < n_sort:      # np.sum(front_id < np.inf), si se cumple la condicion devuelve un 1, en caso contrario devuelve 0
            max_front += 1 #Aumenta uno a la frontera
            for i in np.where(front_id==np.inf)[0]: #si encuentra inf en el front_id
                #if np.sum(front_id < np.inf) >= n_sort: break
                dominated = False #Al principio no hay dominación
                for j in range(i, 0, -1):#Recorrido descendente, desde N hasta 0
                    #print("j--",j)
                    if front_id[j-1] == max_front: 
                        #print("cumple front_id:",front_id)
                        m=2
                        while(m<=self.moo.m) and (sorted_cost[i, m-1] >= sorted_cost[j-1, m-1]): #Compara cada elemento de la segunda funcion objetivo para ver si cumple las condiciones
                            m += 1
                        dominated = m > self.moo.m
                        if dominated or self.moo.m==2:
                            break
                if not dominated:
                    front_id[i] = max_front
        
        return front_id[loc], max_front #Regresa el Valor maximo de frontera y el vector con cada indice de la frontera de los N individuos
        
    def crowding_distance(self, front_id):
        '''
        Calcula la distancia de amontonamiento por cada frontera de Pareto
        Parametro de entrada: front_id, de dimensión (Nx1) donde front_i[i] es el número de frontera de pop[i]
        Valor de retorno: crowd_dis de dimensión (Nx1) contiene las distancias de amontonamiento
        '''
        N = self.pop[0].shape[0] #Numero de individuos en la poblacion
        pop_cost = self.pop[1] #Matriz con individuos de la población evaluados en cada una de las funciones objetivo 
        crowd_dis = np.zeros(N) # Vector de dimension (1xN)
        fronts = np.unique(front_id) #Obtiene las fronteras no repetidas en fron_id
        fronts = fronts[fronts!=np.inf] #Deja todos las fronteras diferentes de "inf"
        
        for f in range(len(fronts)): #Recorre cada una de las N fronteras
            front = np.where(front_id==f+1)[0] #Selecciona elementos(ids) de la frontera si el identificador de la frontera es igual a la frontera (f) actualmente evaluada

            fmax = np.max(pop_cost[front, :], axis=0) #Obtiene los valores maximos de cada funcion objetivo que están en front
            fmin = np.min(pop_cost[front, :], axis=0) #Obtiene los valores minimos de cada funcion objetivo que están en front

            for i in range(self.moo.m): #Ciclo desde 0 hasta (m) funciones objetivo
                rank = np.argsort(pop_cost[front, i]) #Ordena de menor a mayor los indices de cada uno de los valores de la frontera (f)

                crowd_dis[front[rank[0]]] = np.inf #Coloca un "inf" en cada uno de los primeros indices de rank y almacenarlo en crowd_dis

                crowd_dis[front[rank[-1]]] = np.inf #Coloca un "inf" en cada uno de los ultimos indices de rank y almacenarlo en crowd_dis

                for j in range(1, len(front)-1):  #Se recorre los valores de la frontera pero sin considerar el valor del inicio y del final
                    crowd_dis[front[rank[j]]] = crowd_dis[front[rank[j]]] + \
                            (pop_cost[front[rank[j+1]], i] - pop_cost[front[rank[j-1]], i]) / (fmax[i]-fmin[i]) #realiza la siguiente operacion para cada funcion objetivo(j+1 - j-1) /(fmax - fmin)
        return crowd_dis
    
    def tournament(self, fit, K=2):
        '''
        Realiza torneo entre los individuos de la poblacion
        Parametros de entrada:
            K: Número de parámetros a considerar para el fitness
            fit: (N, K) Matriz de valores fitness, donde la columna más alta significa mayor preferencia
        Valor de retorno: indices de los individuos que ganaron el torneo
        '''
        n_total = len(fit) #Numero total de individuos

        a = np.random.randint(n_total, size=self.N) #Vector de N  numeros enteros aleatorios
        b = np.random.randint(n_total, size=(self.N, K)) # Matriz de (N x K) numeros enteros aleatorios

        for i in range(self.N): #Ciclo de 0 a N
            for j in range(K): #Ciclo de 0 a K
                for r in range(fit[0, :].size): #Ciclo de 0 al numero de columnas de (fit)
                    if fit[b[i, j], r] < fit[a[i], r]: #Se eligen valores aleatorios dando prioridad a los valores con menor distancia de amontonamiento
                        a[i] = b[i,j]
        return a
    
    def evolve(self, parents, boundary=None):
        ''' Crea la descendencia de los padres a través de cruce + mutación
        Param:
            parents: (Nxm) matriz de padres para participar en la mutación
            boundary: límites (inferior, superior) de las restricciones
        '''
        dis_c, dis_m = 20, 20
        parents = parents[:(len(parents)//2)*2, :]
        (n, d) = parents.shape
        parent_1, parent_2 = parents[:n//2, :], parents[n//2:, :]
        
        ## CRUZA
        beta = np.empty((n//2, d))
        mu = np.random.random((n//2, d))
        beta[mu <= 0.5]= np.power(2 * mu[mu <= 0.5], 1 / (dis_c + 1))
        beta[mu > 0.5] = np.power(2 * mu[mu > 0.5], -1 / (dis_c + 1))
        beta = beta * ((-1)** np.random.randint(2, size=(n // 2, d)))
        beta[np.random.random((n // 2, d)) < 0.5] = 1
        beta[np.tile(np.random.random((n // 2, 1)) > self.p_c, (1, d))] = 1
        # beta=1 significa sin cruza
        
        offspring = np.vstack(((parent_1 + parent_2) / 2 + beta * (parent_1 - parent_2) / 2,
                                   (parent_1 + parent_2) / 2 - beta * (parent_1 - parent_2) / 2))

        ## Mutacion
        site = np.random.random((n, d)) < self.p_m / d
        mu = np.random.random((n, d))
        temp = site & (mu <= 0.5)
        if boundary is None:
            lower, upper = np.tile(self.moo.lower, (n, 1)), np.tile(self.moo.upper, (n,1))
        else:
            lower, upper = boundary
        norm = (offspring[temp] - lower[temp]) / (upper[temp] - lower[temp])
        offspring[temp] += (upper[temp] - lower[temp]) * \
                               (np.power(2. * mu[temp] + (1. - 2. * mu[temp]) * np.power(1. - norm, dis_m + 1.),
                                         1. / (dis_m + 1)) - 1.)
        temp = ~temp
        norm = (upper[temp] - offspring[temp]) / (upper[temp] - lower[temp])
        offspring[temp] += (upper[temp] - lower[temp])* \
                               (1. - np.power(
                                   2. * (1. - mu[temp]) + 2. * (mu[temp] - 0.5) * np.power(1. - norm, dis_m + 1.),
                                   1. / (dis_m + 1.)))
        offspring = np.maximum(np.minimum(offspring, upper), lower)
        return offspring
    
    def selection(self):
        '''Realiza la selección del entorno de front_ids y crowding_distance
        '''
        front_id, max_front = self.non_dominated_sort(n_sort=self.N)
        next_label = np.zeros(front_id.shape[0], dtype=bool)
        next_label[front_id<max_front] = True
        crowd_dis = self.crowding_distance(front_id)
        last = np.where(front_id==max_front)[0]
        rank = np.argsort(-crowd_dis[last])
        delta_n = rank[:(self.N - int(np.sum(next_label)))]
        next_label[last[delta_n]] = True
        index = np.where(next_label)[0]
        self.pop = [self.pop[0][index, :], self.pop[1][index, :]]
        return front_id[index], crowd_dis[index], index
    
    def run(self):

        #Crea una clasificacion no dominada. #Regresa el Valor maximo de frontera y el vector con cada indice de la frontera de los N individuos
        front_id, max_front = self.non_dominated_sort() 

        #Calcula distancia de amontonamienton de cada elemento por funcion objetivo
        crowd_dis = self.crowding_distance(front_id)
        
        eval_left, n_iter = self.max_evals, 0 #Número máximo de evaluaciones de funciones permitidas (predeterminado = 50,000 por individuo), iteraciones = 0


        while eval_left >= 0 and n_iter <= self.max_iter: #Clico mientras no se alcanza el número máximo de iteraciones
            fit = np.vstack((front_id, crowd_dis)).T #Concatena dos arreglos:  front_id y crowd_dis y los transpone
            print("FIT--",fit)

            mating_pool = self.tournament(fit) #Realiza torneo entre los individuos de la población  de acuerdo con la distancia de amontonamiento

            print("mating_pool--",mating_pool)

            parent = [self.pop[0][mating_pool, :], self.pop[1][mating_pool, :]] #Selecciona a los padres a partir del torneo y la población inicial

            offspring = self.evolve(parent[0]) #Crea la cruza y mutación de los padres

            print("OFFSPRING--",offspring) 

            offspring_cost = self.moo.evaluate(offspring) #Evaluación de la nueva población P+1, mediante las N funciones objetivo

            print("OFFSPRING COST--",offspring_cost) 

            print("POP ANTES--",self.pop)

            self.pop = [np.vstack((self.pop[0], offspring)), np.vstack((self.pop[1], offspring_cost))] #Se crea una nueva población de tamaño 2*N

            print("POP DESPUÉS--",self.pop)

            front_id, crowd_dis, _ = self.selection()
            eval_left -= self.N
            n_iter +=1
                
        print("Número de iteration: {}, Número de evaluación de fuciones: {}".format(n_iter, self.max_evals-eval_left))
        return self.pop
        
    def visualize(self):
        _ = self.run() #Comienza la ejecucion de NSGA una vez que se generó la población inciial y esta fue evaluada

        pop_cost = self.pop[1]
        front_id, max_front = self.non_dominated_sort()
        non_dominated = pop_cost[front_id==1, :]

        print("Valores de X: ",self.pop)
        print("Valores no dominados: ",non_dominated)

        if self.moo.m == 2:
            ax = plt.subplot(111)
            ax.scatter(non_dominated[:, 0], non_dominated[:, 1])

        elif self.moo.m == 3:
            x, y, z = non_dominated[:, 0], non_dominated[:, 1], non_dominated[:, 2]
            ax = plt.subplot(111, projection='3d')
            norm =  (non_dominated - non_dominated.min(0)) / non_dominated.ptp(0)
            ax.scatter(x, y, z, c=norm)

        elif self.moo.m == 4:
            x, y, z,w = non_dominated[:, 0], non_dominated[:, 1], non_dominated[:, 2],non_dominated[:, 3]
            ax = plt.subplot(111, projection='4d')
            norm =  (non_dominated - non_dominated.min(0)) / non_dominated.ptp(0)
            ax.scatter(x, y, z,w, c=norm) 

        else:
            
            for i in range(len(non_dominated)):
                ax = plt.subplot(3,len(non_dominated)//3+1, i)
                ax.plot(range(1, self.moo.m + 1), non_dominated[i, :])
        return ax




################################################################################################
######################PROGRAMA PRINCIPAL########################################################

#FUNCIONES OBJETIVO A MINIMIZAR
#f_1(x) = (x + 2) ** 2 - 10
#f_2(x) = (x - 2) ** 2 + 20

#primera función objetivo a optimizar
def function1(x):
    value = (np.square(x+2) - 10)
    return value

#Segunda función objetivo a optimizar
def function2(x):
    value = (np.square(x-2) + 20)
    return value

#Las funciones objetivo se devuelven como  un objeto y se envian a MOO_Problem para que construya las matrices con las que se va a trabajar
#d: Dimensión del vector de entrada para las funciones objetivo -1
moo_problem = MOO_Problem(function1, function2, d=1)

#Se crea una instancia de la clase NSGA-II, se envia el tamaño de la poblacion (pop_size) y el número de iteraciones (n_iter)
nsgaii = NSGAII(moo_problem, pop_size=10, n_iter=10)

#Se genera la visualización de los datos en forma gráfica
ax = nsgaii.visualize()
ax.set_title('Ejemplo problema MOOP')
ax.set_xlabel('f1(x)')
ax.set_ylabel('f2(x)')
plt.show()

###################### FIN DEL PROGRAMA PRINCIPAL###############################################
################################################################################################
'''
m = 3
d = 4
lb = np.array([[0]*d])
ub = np.array([[1]*d])

def g(x):
    return np.sum(np.square(x[:, 2:]-0.5), axis=1)

def f1(x):
    return -((1+g(x))*np.cos(np.pi/2*x[:, 0])*np.cos(np.pi/2*x[:, 1]))

def f2(x):
    return -((1+g(x))*np.sin(np.pi/2*x[:, 1])*np.cos(np.pi/2*x[:, 0]))

def f3(x):
    return -((1+g(x))*np.sin(np.pi/2*x[:, 0]))

moo_problem = MOO_Problem(f1, f2, f3, d=d, lower=lb, upper=ub)
nsgaii = NSGAII(moo_problem, pop_size=100, n_iter=100)
ax = nsgaii.visualize()
ax.view_init(elev=-145., azim=45)
ax.set_title('DTLZ 2 Problem')
ax.set_xlabel('f1(x)')
ax.set_ylabel('f2(x)')
ax.set_zlabel('f3(x)')
plt.show()'''

input("Esperando...")