"""Entry point for running the evolution strategy on the benchmark function."""

from __future__ import annotations

import numpy as np

from es import EvolutionStrategyConfig, run_evolution_strategy
from functions import axis_parallel_hyperellipsoid, default_bounds


def run_for_dimension(
    dimension: int,
    *,
    results_dir: str,
    max_generations: int = 200,
    seed: int | None = 17,
    save_generations: list[int] | None = None,
    log: bool = False,
):
    x_min, x_max = default_bounds(dimension)
    search_range = x_max - x_min
    initial_sigma = np.full(dimension, 0.15 * search_range[0], dtype=np.float64)

    config = EvolutionStrategyConfig(
        fitness_func=axis_parallel_hyperellipsoid,
        dimension=dimension,
        x_min=x_min,
        x_max=x_max,
        mu=20,
        lambda_=80,
        mutation_probability=0.7,
        initial_sigma=initial_sigma,
        max_generations=max_generations,
        selection="comma",
        recombination="intermediate",
        parents_per_offspring=2,
        success_rule_window=5,
        success_rule_target=0.2,
        sigma_increase=1.22,
        sigma_decrease=0.82,
        sigma_scale_min=1e-3,
        sigma_scale_max=50.0,
        sigma_min=1e-5,
        sigma_max=2.0,
        best_value_threshold=1e-6,
        max_stagnation_generations=40,
        save_generations=save_generations,
        results_dir=results_dir,
        log_every_generation=log,
        seed=seed,
    )

    result = run_evolution_strategy(config)

    print("=" * 80)
    print(f"Результаты для размерности n={dimension}")
    print(f"Лучшее решение: {result.best_generation.best.x}")
    print(f"Лучшее значение функции: {result.best_generation.best.fitness:.8f}")
    print(f"Количество поколений: {result.generations_count}")
    print(f"Время выполнения: {result.time_ms:.2f} мс")
    print("=" * 80)

    return result


def main() -> None:
    # Для n=2 дополнительно сохраняем графики поколений
    run_for_dimension(
        2,
        results_dir="lab5_results_2d",
        save_generations=[1, 2, 3, 5, 8, 10, 15, 20, 25, 30, 40, 50, 75, 100, 150, 200],
        log=True,
    )

    # Для n=3 графики не строим, но выводим статистику
    run_for_dimension(
        3,
        results_dir="lab5_results_3d",
        save_generations=None,
        log=False,
    )


if __name__ == "__main__":
    main()