import math
import os
import random
import shutil
import time
from dataclasses import dataclass
from typing import Callable, List, Tuple

import matplotlib.pyplot as plt
import numpy as np


def target_function(x: float) -> float:
    """f(x) = sin(x)/x^2"""
    return math.sin(x) / (x * x)


def bits_for_precision(x_min: float, x_max: float, digits_after_decimal: int) -> int:
    """
    Подбор числа бит L так, чтобы шаг сетки был ≤ 10^{-digits_after_decimal}.
    Идея как в методичке: 2^L >= (x_max - x_min)*10^{digits_after_decimal}.
    """
    required_levels = int(math.ceil((x_max - x_min) * (10**digits_after_decimal)))
    L = 1
    while (1 << L) < required_levels:
        L += 1
    return L


def decode_bits_to_x(bits: List[int], x_min: float, x_max: float) -> float:
    """Линейное отображение битовой строки в x ∈ [x_min, x_max]."""
    v = 0
    for b in bits:
        v = (v << 1) | int(b)
    L = len(bits)
    levels = (1 << L) - 1
    return x_min + (x_max - x_min) * (v / levels)


def random_bits(L: int) -> List[int]:
    return [random.randint(0, 1) for _ in range(L)]


def eval_population(
    population: List[List[int]],
    x_min: float,
    x_max: float,
    fitness_func: Callable[[float], float],
) -> Tuple[List[float], List[float]]:
    """Оценка популяции: преобразование в x значения и вычисление фитнес функции."""
    xs = [decode_bits_to_x(ch, x_min, x_max) for ch in population]
    fits = [fitness_func(x) for x in xs]
    return xs, fits


def clear_results_directory(results_dir: str) -> None:
    """Очищает папку с результатами перед началом эксперимента."""
    if os.path.exists(results_dir):
        shutil.rmtree(results_dir)
    os.makedirs(results_dir, exist_ok=True)


def reproduction(
    population: List[List[int]], fitnesses: List[float]
) -> List[List[int]]:
    """Репродукция (селекция) методом рулетки."""
    # Чтобы работать с отрицательными f, сдвигаем значения фитнес функции на минимальное
    # значение в популяции. Вычитаем min_fit, т. к. min_fit может быть отрицательным.
    min_fit = min(fitnesses)
    shifted_fitnesses = [f - min_fit + 1e-12 for f in fitnesses]
    s = sum(shifted_fitnesses)

    probs = [sf / s for sf in shifted_fitnesses]
    cum = np.cumsum(probs)
    selected = []
    for _ in population:
        r = random.random()
        idx = int(np.searchsorted(cum, r, side="left"))
        selected.append(population[idx][:])
    return selected


def crossover_pair(
    p1: List[int], p2: List[int], pc: float
) -> Tuple[List[int], List[int]]:
    """Кроссинговер между двумя хромосомами с вероятностью pc."""
    if random.random() <= pc:
        k = random.randint(1, len(p1) - 1)
        return p1[:k] + p2[k:], p2[:k] + p1[k:]
    else:
        return p1[:], p2[:]


def crossover(population: List[List[int]], pc: float) -> List[List[int]]:
    """Оператор кроссинговера (скрещивания) выполняется с заданной вероятностью pc.
    1. Две хромосомы (родители) выбираются случайно из промежуточной популяции.
    2. Случайно выбирается точка скрещивания - число k из диапазона [1,2,…,n-1],
       где n – длина хромосомы (число бит в двоичном коде).
    3. Две новых хромосомы A', B' (потомки) формируются из A и B путем обмена подстрок
       после точки скрещивания с вероятностью pc. Иначе родители добавляются в новую
       популяцию без изменений.

    Если популяция нечетного размера, то последняя хромосома скрещивается со случайной
    другой хромосомой из популяции. В таком случае одна из хромосом может поучаствовать
    в кроссовере дважды.
    """
    # Создаем копию популяции и перемешиваем её для случайного выбора пар
    shuffled_population = population[:]
    random.shuffle(shuffled_population)

    next_population = []
    pop_size = len(shuffled_population)

    for i in range(0, pop_size, 2):
        p1 = shuffled_population[i]
        p2 = shuffled_population[(i + 1) % pop_size]
        c1, c2 = crossover_pair(p1, p2, pc)
        next_population.append(c1)
        next_population.append(c2)

    return next_population[:pop_size]


def mutation(chrom: List[int], pm: float) -> None:
    """Мутация происходит с вероятностью pm.
    1. В хромосоме случайно выбирается k-ая позиция (бит) мутации.
    2. Производится инверсия значения гена в k-й позиции.
    """
    if random.random() <= pm:
        k = random.randint(0, len(chrom) - 1)
        chrom[k] = 1 - chrom[k]


@dataclass
class GARunConfig:
    x_min: float = 3.1
    x_max: float = 20.0
    precision_digits: int = 3  # точность сетки ~0.001
    pop_size: int = 100  # размер популяции
    pc: float = 0.7  # вероятность кроссинговера
    pm: float = 0.01  # вероятность мутации
    max_generations: int = 200  # максимальное количество поколений
    seed: int | None = None  # seed для генератора случайных чисел
    save_generations: list[int] | None = (
        None  # индексы поколений для сохранения графиков
    )
    results_dir: str = "results"  # папка для сохранения графиков
    variance_threshold: float | None = (
        None  # порог дисперсии для остановки (если None - не используется)
    )
    min_fitness_avg: float | None = (
        None  # порог среднего значения фитнес функции для остановки
    )


@dataclass
class GARunResult:
    best_x: float
    best_f: float
    generations: int
    history_best_x: List[float]
    history_best_f: List[float]
    history_populations_x: List[List[float]]  # история всех популяций (x значения)
    history_populations_f: List[List[float]]  # история всех популяций (f значения)
    time_ms: float
    L: int  # число бит


def genetic_algorithm(
    config: GARunConfig,
    fitness_func: Callable[[float], float] = target_function,
) -> GARunResult:
    if config.seed is not None:
        random.seed(config.seed)
        np.random.seed(config.seed)

    # Очищаем папку результатов если нужно сохранять графики
    if config.save_generations:
        clear_results_directory(config.results_dir)

    L = bits_for_precision(config.x_min, config.x_max, config.precision_digits)
    population = [random_bits(L) for _ in range(config.pop_size)]

    start = time.perf_counter()
    history_best_x, history_best_f = [], []
    history_populations_x, history_populations_f = [], []
    best_x, best_f = 0, -float("inf")

    for generation in range(config.max_generations):
        xs, fits = eval_population(population, config.x_min, config.x_max, fitness_func)

        # лучший в поколении + глобально лучший
        gi = int(np.argmax(fits))
        gen_best_x, gen_best_f = xs[gi], fits[gi]
        if gen_best_f > best_f:
            best_x, best_f = gen_best_x, gen_best_f

        history_best_x.append(gen_best_x)
        history_best_f.append(gen_best_f)
        history_populations_x.append(xs[:])
        history_populations_f.append(fits[:])

        # Проверка критериев остановки
        stop_algorithm = False

        # Критерий остановки по дисперсии
        if config.variance_threshold is not None:
            fitness_variance = np.var(fits)
            if fitness_variance < config.variance_threshold:
                print(
                    f"Остановка на поколении {generation}: дисперсия {fitness_variance:.6f} < {config.variance_threshold}"
                )
                stop_algorithm = True

        # Критерий остановки по среднему значению фитнес функции
        if config.min_fitness_avg is not None:
            mean_fitness = np.mean(fits)
            if mean_fitness > config.min_fitness_avg:
                print(
                    f"Остановка на поколении {generation}: среднее значение {mean_fitness:.6f} > {config.min_fitness_avg}"
                )
                stop_algorithm = True

        # Сохраняем график последнего поколения при досрочной остановке
        if stop_algorithm:
            if config.save_generations:
                plot_generation_snapshot(
                    history_best_x,
                    history_best_f,
                    history_populations_x,
                    history_populations_f,
                    generation,
                    config.x_min,
                    config.x_max,
                    config.results_dir,
                )
            break

        # Сохранение графика для указанных поколений
        if config.save_generations and generation in config.save_generations:
            plot_generation_snapshot(
                history_best_x,
                history_best_f,
                history_populations_x,
                history_populations_f,
                generation,
                config.x_min,
                config.x_max,
                config.results_dir,
            )

        # селекция
        parents = reproduction(population, fits)

        # кроссинговер попарно
        next_population = crossover(parents, config.pc)

        # мутация
        for ch in next_population:
            mutation(ch, config.pm)

        population = next_population[: config.pop_size]

    end = time.perf_counter()

    return GARunResult(
        best_x,
        best_f,
        len(history_best_x),  # реальное количество поколений
        history_best_x,
        history_best_f,
        history_populations_x,
        history_populations_f,
        (end - start) * 1000.0,
        L,
    )


def plot_generation_snapshot(
    history_best_x: List[float],
    history_best_f: List[float],
    history_populations_x: List[List[float]],
    history_populations_f: List[List[float]],
    generation: int,
    x_min: float,
    x_max: float,
    results_dir: str,
) -> str:
    """
    График для конкретного поколения с отображением всей популяции.
    """
    os.makedirs(results_dir, exist_ok=True)

    xs = np.linspace(x_min, x_max, 1500)
    ys = np.sin(xs) / (xs * xs)

    fig = plt.figure(figsize=(10, 6))
    plt.plot(xs, ys, label="f(x)=sin(x)/x^2", alpha=0.7, color="blue")

    # Отображаем всю популяцию текущего поколения
    if generation < len(history_populations_x):
        current_pop_x = history_populations_x[generation]
        current_pop_f = history_populations_f[generation]

        # Вся популяция серыми точками
        plt.scatter(
            current_pop_x,
            current_pop_f,
            s=20,
            alpha=0.9,
            color="gray",
            label=f"Популяция поколения {generation}",
        )

        # Лучшая особь красной точкой
        best_idx = np.argmax(current_pop_f)
        plt.scatter(
            [current_pop_x[best_idx]],
            [current_pop_f[best_idx]],
            s=60,
            color="red",
            marker="o",
            label=f"Лучшая особь поколения {generation}",
            edgecolors="darkred",
            linewidth=1,
        )

    # История лучших по поколениям (до текущего включительно)
    if generation > 0:
        history_x_until_now = history_best_x[:generation]
        history_f_until_now = history_best_f[:generation]
        plt.scatter(
            history_x_until_now,
            history_f_until_now,
            s=15,
            alpha=0.9,
            color="orange",
            label="Лучшие предыдущих поколений",
        )

    plt.xlabel("x")
    plt.ylabel("f(x)")
    plt.title(f"Популяция поколения {generation}")
    plt.legend(loc="best")
    plt.grid(True, alpha=0.3)
    plt.tight_layout()

    filename = f"generation_{generation:03d}.png"
    path_png = os.path.join(results_dir, filename)
    plt.savefig(path_png, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return path_png