import random
from math import log

import numpy as np
from numpy.typing import NDArray

from gp import Chromosome
from gp.crossovers import crossover_subtree
from gp.fitness import (
    HuberFitness,
    MAEFitness,
    MSEFitness,
    NRMSEFitness,
    PenalizedFitness,
    RMSEFitness,
)
from gp.ga import GARunConfig, genetic_algorithm
from gp.mutations import (
    grow_mutation,
    hoist_mutation,
    node_replacement_mutation,
    shrink_mutation,
)
from gp.ops import ADD, COS, DIV, EXP, MUL, NEG, POW, SIN, SQUARE, SUB
from gp.population import ramped_initialization
from gp.primitive import Const, Var
from gp.selection import roulette_selection, tournament_selection

NUM_VARS = 9
TEST_POINTS = 10000
MAX_DEPTH = 15
MAX_GENERATIONS = 200
np.random.seed(17)
random.seed(17)
X = np.random.uniform(-5.536, 5.536, size=(TEST_POINTS, NUM_VARS))
# axes = [np.linspace(-5.536, 5.536, TEST_POINTS) for _ in range(NUM_VARS)]
# X = np.array(np.meshgrid(*axes)).T.reshape(-1, NUM_VARS)
operations = [SQUARE, SIN, COS, EXP, ADD, SUB, MUL, DIV, POW]
# operations = [SQUARE, ADD, SUB, MUL]
terminals = [Var(f"x{i}") for i in range(1, NUM_VARS + 1)]


# def target_function(x: NDArray[np.float64]) -> NDArray[np.float64]:
#     """
#     f(x) = x1 + x2 + sin(x1)
#     x имеет форму (n_samples, n_vars)
#     """
#     x1 = x[:, 0]
#     x2 = x[:, 1]
#     return x1 + x2 + np.sin(x1) + np.sin(x2) + np.exp(-x1) + np.cos(x2)


# def target_function(x: NDArray[np.float64]) -> NDArray[np.float64]:
#     """
#     Простая тестовая функция: сумма косинусов всех переменных.
#     f(x) = sum_i cos(x_i)
#     x имеет форму (n_samples, n_vars)
#     """
#     return np.sum(x, axis=1)


def target_function(x: NDArray[np.float64]) -> NDArray[np.float64]:
    """
    Векторизованная версия функции:
    f(x) = sum_{i=1}^n sum_{j=1}^i x_j^2
    x имеет форму (n_samples, n_vars)
    """
    # Префиксные суммы квадратов по оси переменных
    x_sq = x**2
    prefix_sums = np.cumsum(x_sq, axis=1)
    # Суммируем по i (ось 1)
    return np.sum(prefix_sums, axis=1)


# def target_function(x: NDArray[np.float64]) -> NDArray[np.float64]:
#     """
#     Rastrigin function.
#     f(x) = 10 * n + sum(x_i^2 - 10 * cos(2πx_i))
#     x: shape (n_samples, n_vars)
#     """
#     n = x.shape[1]
#     return 10 * n + np.sum(x**2 - 10 * np.cos(2 * np.pi * x), axis=1)

# fitness_function = MSEFitness(target_function, lambda: X)
# fitness = HuberFitness(target_function, lambda: X, delta=1.0)
# fitness_function = PenalizedFitness(
#     target_function, lambda: X, base_fitness=fitness, lambda_=0.003
# )
fitness_function = RMSEFitness(target_function, lambda: X)
# fitness_function = PenalizedFitness(
#     target_function, lambda: X, base_fitness=fitness, lambda_=0.003
# )


def adaptive_mutation(
    chromosome: Chromosome,
    generation: int,
    max_generations: int,
    max_depth: int,
) -> Chromosome:
    """Адаптивная мутация.

    Меняет вероятность типов мутации по ходу эволюции:
    - Ранняя фаза (<30%): 70% grow, 30% shrink
    - Средняя фаза (30–70%): 40% grow, 60% shrink
    - Поздняя фаза (>=70%): 20% grow, 80% shrink
    """

    r = random.random()

    if r < 0.4:
        return grow_mutation(chromosome, max_depth=max_depth)
    elif r < 0.7:
        return node_replacement_mutation(chromosome)
    elif r < 0.85:
        return hoist_mutation(chromosome)

    return shrink_mutation(chromosome)


# def adaptive_mutation(
#     chromosome: Chromosome,
#     generation: int,
#     max_generations: int,
#     max_depth: int,
# ) -> Chromosome:
#     """Адаптивная мутация.

#     Меняет вероятность типов мутации по ходу эволюции:
#     - Ранняя фаза (<30%): 70% grow, 30% shrink
#     - Средняя фаза (30–70%): 40% grow, 60% shrink
#     - Поздняя фаза (>=70%): 20% grow, 80% shrink
#     """

#     # Вычисляем прогресс в диапазоне [0, 1]
#     if max_generations <= 0:
#         progress = 0.0
#     else:
#         progress = min(1.0, max(0.0, generation / max_generations))

#     r = random.random()

#     # Определяем тип мутации
#     if progress < 0.3:
#         do_grow = r < 0.7
#     elif progress < 0.7:
#         do_grow = r < 0.4
#     else:
#         do_grow = r < 0.2

#     # Выполняем выбранную мутацию
#     if do_grow:
#         return grow_mutation(chromosome, max_depth=max_depth)
#     return shrink_mutation(chromosome)


config = GARunConfig(
    fitness_func=fitness_function,
    crossover_fn=lambda p1, p2: crossover_subtree(p1, p2, max_depth=8),
    mutation_fn=lambda chrom, gen_num: adaptive_mutation(
        chrom, gen_num, MAX_GENERATIONS, MAX_DEPTH
    ),
    # selection_fn=roulette_selection,
    selection_fn=lambda p, f: tournament_selection(p, f, k=3),
    init_population=ramped_initialization(
        10, [4, 5, 6, 6, 7, 7, 8, 9, 10, 11], terminals, operations
    ),
    seed=17,
    pc=0.9,
    pm=0.3,
    elitism=10,
    max_generations=MAX_GENERATIONS,
    log_every_generation=True,
)

result = genetic_algorithm(config)


# Выводим результаты
print(f"Лучшая особь: {result.best_generation.best}")
print(result.best_generation.best.root.to_str_tree())
print(f"Лучшее значение фитнеса: {result.best_generation.best_fitness:.6f}")
print(f"Количество поколений: {result.generations_count}")
print(f"Время выполнения: {result.time_ms:.2f} мс")

mse_fitness = MSEFitness(target_function, lambda: X)
print(f"MSE: {mse_fitness(result.best_generation.best):.6f}")
rmse_fitness = RMSEFitness(target_function, lambda: X)
print(f"RMSE: {rmse_fitness(result.best_generation.best):.6f}")
mae_fitness = MAEFitness(target_function, lambda: X)
print(f"MAE: {mae_fitness(result.best_generation.best):.6f}")
huber_fitness = HuberFitness(target_function, lambda: X, delta=1.0)
print(f"Huber: {huber_fitness(result.best_generation.best):.6f}")
nrmse_fitness = NRMSEFitness(target_function, lambda: X)
print(f"NRMSE: {nrmse_fitness(result.best_generation.best):.6f}")