import seaborn as sns
import matplotlib.pyplot as plt
sns.set_theme(style="whitegrid")

import metaheuristic_designer as mhd
from metaheuristic_designer.benchmarks import Rastrigin
from metaheuristic_designer.initializers import UniformInitializer
from metaheuristic_designer.strategies import DE
from metaheuristic_designer.algorithms import Algorithm
from metaheuristic_designer.history_tracker import HistoryTracker
from metaheuristic_designer.stopping_condition import StoppingCondition
from metaheuristic_designer.parameter_schedules import ExponentialDecaySchedule

rng = mhd.check_random_state(42)
DIM = 5
objfunc = Rastrigin(DIM, mode="min")

strategy = DE(
   initializer=UniformInitializer(
      objfunc.dimension, objfunc.lower_bound, objfunc.upper_bound,
      population_size=100, random_state=rng
   ),
   de_operator_name="DE/rand/1",
   F=ExponentialDecaySchedule(init_value=1, final_value=0.05, alpha=0.99),
   Cr=0.9,
   name="DE",
   random_state=rng,
)

algo = mhd.Algorithm(
    objfunc,
    strategy,
    stop_cond="max_iterations",
    max_iterations=200,
    reporter="silent",
    history_tracker=mhd.HistoryTracker(
      track_median=True,
      track_worst=True,
      track_full_objective=True,
      track_diversity=True,
      track_parameters=True,
   )
)
algo.optimize()

df = algo.history_tracker.to_pandas()
full_obj_df = algo.history_tracker.to_pandas_full_objective()

fig, ax = plt.subplots(figsize=(8, 5))
sns.lineplot(data=df, x="iteration", y="best_objective", label="Best", linewidth=2, ax=ax)
sns.lineplot(data=df, x="iteration", y="median_objective", label="Median", linewidth=1.5, ax=ax)
sns.lineplot(data=df, x="iteration", y="worst_objective", label="Worst", linewidth=1, ax=ax)
ax.set_xlabel("Generation")
ax.set_ylabel("Objective")
ax.set_title("Best, Median, and Worst Objectives")
ax.legend()
plt.tight_layout()
plt.show()