import random

import matplotlib.pyplot as plt  # type: ignore
from tsp import TSP
from multi_tsp_plotter import MultiTSPPlotter

random.seed(1)
tsp = TSP([(random.random(), random.random()) for _ in range(50)])
# tsp = TSP.from_file("tsp100.txt")
mu = 5
lam = 20
k = 3
plotter = MultiTSPPlotter(mu)

best_sol = None
best_score = None
population = [tsp.random_tour() for _ in range(mu)]
# population = [tsp.random_grasp_tour(0.2) for _ in range(mu)]
# sol = tsp.random_grasp_tour(0.75)

while True:
    for sol in population:
        if best_sol is None or tsp.score(sol) < best_score:
            best_sol = sol
            best_score = tsp.score(sol)

    tweaks = []
    for sol in population:
        for _ in range(int(lam / mu)):
            tweaks.append(tsp.k_opt(sol, k))
            # print(tsp.score(sol), "->", tsp.score(tweaks[-1]))

    population = sorted(set(tweaks + population), key=lambda sol: tsp.score(sol))
    population = population[:mu]

    # print([tsp.score(sol) for sol in population])

    for index in range(mu):
        plotter.show_tour_in_spot(population[index], index)
    # plotter.show_tour(best_sol)
    plotter.add_to_score_line(best_score)
    plotter.update_best(best_score)
    plt.pause(0.0001)