#!/usr/bin/env python3
"""
DTLZ Multi-Objective Optimization Example

This example demonstrates using the DTLZ1 problem for multi-objective optimization.
We define a fitness function based on the DTLZ1 problem and use a genetic algorithm to
evolve solutions. The results are visualized in a 3D scatter plot.
"""

import matplotlib.pyplot as plt  # type: ignore
import radiate as rd
import numpy as np  # type: ignore
from numba import jit, float64  # type: ignore

rd.random.seed(501)

variables = 4
objectives = 3
k = variables - objectives + 1


@jit(float64[:](float64[:]), nopython=True)
def dtlz_1(val: np.ndarray) -> np.ndarray:
    g_vals = val[variables - k :] - 0.5
    g = 100.0 * (k + np.sum(g_vals**2 - np.cos(20.0 * np.pi * g_vals)))

    base = 0.5 * (1.0 + g)

    f = np.full(objectives, base, dtype=np.float64)

    for i in range(objectives):
        prod_end = objectives - 1 - i
        if prod_end > 0:
            f[i] *= np.prod(val[:prod_end])

        if i > 0:
            f[i] *= 1.0 - val[objectives - 1 - i]

    return f


engine = (
    rd.Engine.float(variables, use_numpy=True)
    .fitness(dtlz_1)
    .objective(rd.MIN, rd.MIN, rd.MIN)
    .front_range(100, 150)
    .select(rd.Select.tournament(k=5), rd.Select.nsga3(points=12))
    .alters(
        rd.Cross.sbx(1.0, 2.0),
        rd.Mutate.uniform(0.1),
    )
    .limit(rd.Limit.generations(2000))
)

result = engine.run(ui=True)
print(result.metrics().dashboard())

front = result.front()

x = [member.score()[0] for member in front]
y = [member.score()[1] for member in front]
z = [member.score()[2] for member in front]

fig = plt.figure()
ax = plt.axes(projection="3d")
ax.scatter(x, y, z)
ax.set_xlim([0.0, 0.5])  # type: ignore
ax.set_ylim([0.0, 0.5])  # type: ignore
ax.set_zlim([0.0, 0.5])  # type: ignore
plt.show()
