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()

long_df = full_obj_df.melt(id_vars="iteration", var_name="individual", value_name="objective")
# Keep every 10th generation to avoid overplotting
plot_data = long_df[long_df["iteration"] % 5 == 0]

fig, ax = plt.subplots(figsize=(12, 5))
sns.boxplot(data=plot_data, x="iteration", y="objective", hue="iteration", ax=ax,
            palette="viridis", width=0.6, legend=False)
ax.set_xlabel("Generation")
ax.set_ylabel("Objective")
ax.set_title("Fitness Distribution Every 5 Generations")
ax.tick_params(axis='x', rotation=-45)
ax.axhline(0, color="grey")
ax.grid()
plt.tight_layout()
plt.show()