import math
import os
import shutil
import statistics

import numpy as np
from gen import GARunConfig, genetic_algorithm
from prettytable import PrettyTable


def fitness_function(chromosome: np.ndarray) -> np.ndarray:
    return chromosome[0] ** 2 + 2 * chromosome[1] ** 2


# Базовая папка для экспериментов
BASE_DIR = "experiments"

# Параметры для экспериментов
POPULATION_SIZES = [10, 25, 50, 100]
PC_VALUES = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8]  # вероятности кроссинговера
PM_VALUES = [0.001, 0.01, 0.05, 0.1, 0.2]  # вероятности мутации

# Количество запусков для усреднения результатов
NUM_RUNS = 1

# Базовые параметры (как в main.py)
BASE_CONFIG = {
    "x_min": np.array([-5.12, -5.12]),
    "x_max": np.array([5.12, 5.12]),
    "fitness_func": fitness_function,
    "max_generations": 200,
    "seed": None,  # None для случайности, т. к. всё усредняем
    "minimize": True,
    "fitness_avg_threshold": 0.05,  # критерий остановки
    # при включенном сохранении графиков на время смотреть бессмысленно
    # "save_generations": [1, 50, 199],
}


def run_single_experiment(
    pop_size: int, pc: float, pm: float
) -> tuple[float, float, float, float]:
    """
    Запускает несколько экспериментов с заданными параметрами и усредняет результаты.
    Возвращает (среднее_время_в_мс, стд_отклонение_времени, среднее_поколений, стд_отклонение_поколений).
    """
    times = []
    generations = []

    for run_num in range(NUM_RUNS):
        config = GARunConfig(
            **BASE_CONFIG,
            pop_size=pop_size,
            pc=pc,
            pm=pm,
            results_dir=os.path.join(
                BASE_DIR,
                str(pop_size),
                f"pc_{pc:.3f}",
                f"pm_{pm:.3f}",
                f"run_{run_num}",
            ),
        )

        result = genetic_algorithm(config)
        times.append(result.time_ms)
        generations.append(result.generations_count)

    # Вычисляем средние значения и стандартные отклонения
    avg_time = statistics.mean(times)
    std_time = statistics.stdev(times) if len(times) > 1 else 0.0
    avg_generations = statistics.mean(generations)
    std_generations = statistics.stdev(generations) if len(generations) > 1 else 0.0

    return avg_time, std_time, avg_generations, std_generations


def run_experiments_for_population(pop_size: int) -> PrettyTable:
    """
    Запускает эксперименты для одного размера популяции.
    Возвращает таблицу результатов.
    """
    print(f"\nЗапуск экспериментов для популяции размером {pop_size}...")
    print(f"Количество запусков для усреднения: {NUM_RUNS}")

    # Создаем таблицу
    table = PrettyTable()
    table.field_names = ["Pc \\ Pm"] + [f"{pm:.3f}" for pm in PM_VALUES]

    # Запускаем эксперименты для всех комбинаций Pc и Pm
    for pc in PC_VALUES:
        row = [f"{pc:.1f}"]
        for pm in PM_VALUES:
            print(f"  Эксперимент: pop_size={pop_size}, Pc={pc:.1f}, Pm={pm:.3f}")
            avg_time, std_time, avg_generations, std_generations = (
                run_single_experiment(pop_size, pc, pm)
            )

            # Форматируем результат: среднее_время±стд_отклонение (среднее_поколения±стд_отклонение)
            # cell_value = f"{avg_time:.1f}±{std_time:.1f} ({avg_generations:.1f}±{std_generations:.1f})"
            cell_value = f"{avg_time:.1f} ({avg_generations:.0f})"

            if avg_generations == BASE_CONFIG["max_generations"]:
                cell_value = "—"

            row.append(cell_value)
        table.add_row(row)

    return table


def main():
    """Основная функция для запуска всех экспериментов."""
    print("=" * 60)
    print("ЗАПУСК ЭКСПЕРИМЕНТОВ ПО ПАРАМЕТРАМ ГЕНЕТИЧЕСКОГО АЛГОРИТМА")
    print("=" * 60)
    print(f"Размеры популяции: {POPULATION_SIZES}")
    print(f"Значения Pc: {PC_VALUES}")
    print(f"Значения Pm: {PM_VALUES}")
    print(f"Количество запусков для усреднения: {NUM_RUNS}")
    print(
        f"Критерий остановки: среднее значение > {BASE_CONFIG['fitness_avg_threshold']}"
    )
    print("=" * 60)

    # Создаем базовую папку
    if os.path.exists(BASE_DIR):
        shutil.rmtree(BASE_DIR)
    os.makedirs(BASE_DIR)

    # Запускаем эксперименты для каждого размера популяции
    for pop_size in POPULATION_SIZES:
        table = run_experiments_for_population(pop_size)

        print(f"\n{'='*60}")
        print(f"РЕЗУЛЬТАТЫ ДЛЯ ПОПУЛЯЦИИ РАЗМЕРОМ {pop_size}")
        print(f"{'='*60}")
        print(
            f"Формат: среднее_время±стд_отклонение_мс (среднее_поколения±стд_отклонение)"
        )
        print(f"Усреднено по {NUM_RUNS} запускам")
        print(table)

        pop_exp_dir = os.path.join(BASE_DIR, str(pop_size))
        os.makedirs(pop_exp_dir, exist_ok=True)
        with open(os.path.join(pop_exp_dir, "results.csv"), "w", encoding="utf-8") as f:
            f.write(table.get_csv_string())
        print(f"Результаты сохранены в папке: {pop_exp_dir}")

    print(f"\n{'='*60}")
    print("ВСЕ ЭКСПЕРИМЕНТЫ ЗАВЕРШЕНЫ!")
    print(f"Результаты сохранены в {BASE_DIR}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()