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

param_cols = ["DE/rand/1.F", "DE/rand/1.Cr"]

fig, ax = plt.subplots(figsize=(8, 6))
for col in param_cols:
   sns.lineplot(data=df, x="iteration", y=col, ax=ax, label=col)
ax.set_xlabel("Generation")
ax.set_ylabel("Parameter value")
ax.set_title("Scheduled Parameter Evolution")
ax.legend()
ax.axvline(0, color="grey")
ax.axhline(0, color="grey")
plt.tight_layout()
plt.show()