from cProfile import label
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.algorithms.moo.ctaea import CTAEA
from pymoo.algorithms.moo.nsga3 import NSGA3

from pymoo.factory import  get_reference_directions
from pymoo.optimize import minimize

from pymoo.visualization.scatter import Scatter
from reproblem import *

import time

# Define the problem
problem = CRE25() # CRE21() CRE22() CRE23() CRE24() CRE25()

ref_dirs = get_reference_directions("das-dennis", 2, n_partitions=64)

# Define Algorithms
nsga_2_alg = NSGA2(pop_size=100)
nsga_3_alg = NSGA3(pop_size=100, ref_dirs=ref_dirs)
ctaea_alg  = CTAEA(ref_dirs=ref_dirs)

# C-TAEA
start = time.time()
res_ctaea = minimize(problem,
               ctaea_alg,
               ('n_gen', 200),
               seed=1,
               verbose=False)
end = time.time()
print('CTAEA', end - start)

# NSGA-II
res_nsga_2 = minimize(problem,
               nsga_2_alg,
               ('n_gen', 200),
               seed=1,
               verbose=False)
end = time.time()
print('NSGA2', end - start)

# NSGA-III
res_nsga_3 = minimize(problem,
               nsga_3_alg,
               ('n_gen', 200),
               seed=1,
               verbose=False)
end = time.time()
print('NSGA3', end - start)

# Plot the results
plot = Scatter(title="Approximated Pareto fronts of the CRE2-4-5", legend=True)

plot.add(res_ctaea.F, facecolor="none", edgecolor="blue", label="C-TAEA")
plot.add(res_nsga_2.F, facecolor="none", edgecolor="red", label="NSGA-II")
plot.add(res_nsga_3.F, facecolor="none", edgecolor="yellow", label="NSGA-III")

plot.show()