"""
graphviz должен быть доступен в PATH (недостаточно просто установить через pip)

Можно проверить командой
dot -V
"""

import random

import numpy as np
from numpy.typing import NDArray

from gp.crossovers import crossover_subtree
from gp.fitness import (
    MAEFitness,
    MSEFitness,
    NRMSEFitness,
    RMSEFitness,
)
from gp.ga import GARunConfig, genetic_algorithm
from gp.mutations import (
    CombinedMutation,
    GrowMutation,
    HoistMutation,
    NodeReplacementMutation,
    ShrinkMutation,
)
from gp.ops import ADD, COS, DIV, EXP, MUL, POW, SIN, SQUARE, SUB
from gp.primitive import Var
from gp.selection import tournament_selection
from gp.utils import ramped_initialization

NUM_VARS = 8
TEST_POINTS = 10000
MAX_DEPTH = 10
MAX_GENERATIONS = 200
SEED = 17
np.random.seed(SEED)
random.seed(SEED)
X = np.random.uniform(-5.536, 5.536, size=(TEST_POINTS, NUM_VARS))
operations = [SQUARE, SIN, COS, EXP, ADD, SUB, MUL, DIV, POW]
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) = 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)


fitness_function = RMSEFitness(target_function, lambda: X)
combined_mutation = CombinedMutation(
    mutations=[
        GrowMutation(max_depth=MAX_DEPTH),
        NodeReplacementMutation(),
        HoistMutation(),
        ShrinkMutation(),
    ],
    probs=[0.4, 0.3, 0.15, 0.15],
)

init_population = ramped_initialization(
    20, [i for i in range(MAX_DEPTH - 9, MAX_DEPTH + 1)], terminals, operations
)

print("Population size:", len(init_population))

config = GARunConfig(
    fitness_func=fitness_function,
    crossover_fn=lambda p1, p2: crossover_subtree(p1, p2, max_depth=MAX_DEPTH),
    mutation_fn=combined_mutation,
    selection_fn=lambda p, f: tournament_selection(p, f, k=3),
    init_population=init_population,
    seed=SEED,
    pc=0.85,
    pm=0.15,
    elitism=15,
    max_generations=MAX_GENERATIONS,
    log_every_generation=True,
    save_generations=[1, 10, 20, 30, 40, 50, 100, 150, 200],
)

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} мс")
print("Population size:", len(init_population))

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}")
nrmse_fitness = NRMSEFitness(target_function, lambda: X)
print(f"NRMSE: {nrmse_fitness(result.best_generation.best):.6f}")