import math
import os

import matplotlib.pyplot as plt
import numpy as np
from gen import (
    Chromosome,
    GARunConfig,
    genetic_algorithm,
    initialize_random_population,
    inversion_mutation_fn,
    partially_mapped_crossover_fn,
    plot_fitness_history,
    plot_tour,
    swap_mutation_fn,
)

# В списке из 89 городов только 38 уникальных
cities = set()
with open("data.txt", "r") as file:
    for line in file:
        # x и y поменяны местами в визуализациях в методичке
        _, y, x = line.split()
        cities.add((float(x), float(y)))
cities = list(cities)


def euclidean_distance(city1, city2):
    return math.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2)


def build_fitness_function(cities):
    def fitness_function(chromosome: Chromosome) -> float:
        return sum(
            euclidean_distance(cities[chromosome[i]], cities[chromosome[i + 1]])
            for i in range(len(chromosome) - 1)
        ) + euclidean_distance(cities[chromosome[0]], cities[chromosome[-1]])

    return fitness_function


config = GARunConfig(
    fitness_func=build_fitness_function(cities),
    initialize_population_fn=initialize_random_population,
    cities=cities,
    crossover_fn=partially_mapped_crossover_fn,
    # mutation_fn=swap_mutation_fn,
    mutation_fn=inversion_mutation_fn,
    pop_size=500,
    elitism=3,
    pc=0.9,
    pm=0.3,
    max_generations=2500,
    # max_best_repetitions=10,
    minimize=False,
    seed=17,
    save_generations=[
        1,
        5,
        20,
        50,
        100,
        300,
        500,
        700,
        900,
        1500,
        2000,
        2500,
        3000,
        3500,
        4000,
        4500,
    ],
    log_every_generation=True,
)

result = genetic_algorithm(config)

# Сохраняем конфиг и результаты в файлы
config.save()
result.save(config.results_dir)

# Выводим результаты
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} мс")

# Сохраняем лучшую особь за всё время
fig = plt.figure(figsize=(7, 7))
fig.suptitle(
    f"Поколение #{result.best_generation.number}. "
    f"Лучшая особь: {result.best_generation.best_fitness:.4f}. "
    f"Среднее значение: {np.mean(result.best_generation.fitnesses):.4f}",
    fontsize=14,
    y=0.95,
)

# Рисуем лучший маршрут в поколении
ax = fig.add_subplot(1, 1, 1)
plot_tour(config.cities, result.best_generation.best, ax)
filename = f"best_generation_{result.best_generation.number:03d}.png"
path_png = os.path.join(config.results_dir, filename)
fig.savefig(path_png, dpi=150, bbox_inches="tight")
plt.close(fig)

# Рисуем график прогресса по поколениям
plot_fitness_history(
    result, save_path=os.path.join(config.results_dir, "fitness_history.png")
)