import matplotlib.pyplot as plt
import numpy as np
from gen import GARunConfig, genetic_algorithm


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


config = GARunConfig(
    x_min=np.array([-5.12, -5.12]),
    x_max=np.array([5.12, 5.12]),
    fitness_func=fitness_function,
    pop_size=25,
    pc=0.5,
    pm=0.01,
    max_generations=200,
    max_best_repetitions=10,
    minimize=True,
    seed=17,
    save_generations=[1, 2, 3, 5, 7, 9, 10, 15, 19],
    log_every_generation=True,
)

result = genetic_algorithm(config)

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