"""Parameter sweep experiments for the evolution strategy."""

from __future__ import annotations

import statistics
from pathlib import Path
from typing import Iterable

import numpy as np
from prettytable import PrettyTable

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

POPULATION_SIZES = [5, 10, 20, 40]
MUTATION_PROBABILITIES = [0.3, 0.5, 0.7, 0.9, 1.0]
NUM_RUNS = 5
LAMBDA_FACTOR = 5
RESULTS_DIR = Path("lab5_experiments")


def build_config(dimension: int, mu: int, mutation_probability: float) -> EvolutionStrategyConfig:
    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)
    return EvolutionStrategyConfig(
        fitness_func=axis_parallel_hyperellipsoid,
        dimension=dimension,
        x_min=x_min,
        x_max=x_max,
        mu=mu,
        lambda_=mu * LAMBDA_FACTOR,
        mutation_probability=mutation_probability,
        initial_sigma=initial_sigma,
        max_generations=300,
        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=80,
        save_generations=None,
        results_dir=str(RESULTS_DIR / "tmp"),
        log_every_generation=False,
        seed=None,
    )


def run_single_experiment(config: EvolutionStrategyConfig) -> tuple[float, int, float]:
    result = run_evolution_strategy(config)
    return result.time_ms, result.generations_count, result.best_generation.best.fitness


def summarize(values: Iterable[float]) -> tuple[float, float]:
    values = list(values)
    if not values:
        return 0.0, 0.0
    if len(values) == 1:
        return values[0], 0.0
    return statistics.mean(values), statistics.stdev(values)


def run_grid_for_dimension(dimension: int) -> PrettyTable:
    table = PrettyTable()
    table.field_names = ["mu \\ p_mut"] + [f"{pm:.2f}" for pm in MUTATION_PROBABILITIES]

    for mu in POPULATION_SIZES:
        row = [str(mu)]
        for pm in MUTATION_PROBABILITIES:
            times: list[float] = []
            generations: list[int] = []
            best_values: list[float] = []

            for run_idx in range(NUM_RUNS):
                config = build_config(dimension, mu, pm)
                # Для воспроизводимости меняем seed для каждого запуска
                config.seed = np.random.randint(0, 1_000_000)
                time_ms, gens, best = run_single_experiment(config)
                times.append(time_ms)
                generations.append(gens)
                best_values.append(best)

            avg_time, std_time = summarize(times)
            avg_gen, std_gen = summarize(generations)
            avg_best, std_best = summarize(best_values)

            cell = f"{avg_time:.1f}±{std_time:.1f} ({avg_gen:.0f}±{std_gen:.0f}) {avg_best:.4f}"
            row.append(cell)
        table.add_row(row)

    return table


def save_table(table: PrettyTable, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        f.write(table.get_csv_string())


def main() -> None:
    if RESULTS_DIR.exists():
        for child in RESULTS_DIR.iterdir():
            if child.is_file():
                child.unlink()

    print("=" * 80)
    print("Исследование параметров эволюционной стратегии")
    print("Популяции:", POPULATION_SIZES)
    print("Вероятности мутации:", MUTATION_PROBABILITIES)
    print(f"Каждая конфигурация запускается {NUM_RUNS} раз")
    print("=" * 80)

    for dimension in (2, 3):
        print(f"\nРезультаты для размерности n={dimension}")
        table = run_grid_for_dimension(dimension)
        print(table)
        save_table(table, RESULTS_DIR / f"dimension_{dimension}.csv")
        print(f"Таблица сохранена в {RESULTS_DIR / f'dimension_{dimension}.csv'}")


if __name__ == "__main__":
    main()