import matplotlib.pyplot as plt
from libs.clustering import split_tour_across_clusters
from libs.aco import AntColony, total_distance
import numpy as np

cities = [[6734, 1453], [2233, 10], [5530, 1424], [401, 841], [3082, 1644], [7608, 4458], [7573, 3716], [7265, 1268], [6898, 1885], [1112, 2049], [5468, 2606], [5989, 2873], [4706, 2674], [4612, 2035], [6347, 2683], [6107, 669], [7611, 5184], [7462, 3590], [7732, 4723], [5900, 3561], [4483, 3369], [6101, 1110], [5199, 2182], [1633, 2809], [4307, 2322], [675, 1006], [7555, 4819], [7541, 3981], [3177, 756], [7352, 4506], [7545, 2801], [3245, 3305], [6426, 3173], [4608, 1198], [23, 2216], [7248, 3779], [7762, 4595], [7392, 2244], [3484, 2829], [6271, 2135], [4985, 140], [1916, 1569], [7280, 4899], [7509, 3239], [10, 2676], [6807, 2993], [5185, 3258], [3023, 1942]]
optimal = 33523

max_times = [0.1, 0.2, 0.5, 1, 2]
n_ants = 10
alpha = 1
betas = [2, 4, 6, 8]  # Pour contenir deux valeurs de beta
evaporation = 0.5
intensification = 2
n_runs = 2  # Nombre de fois où chaque configuration sera exécutée pour obtenir une moyenne

average_best_route_per_max_time_per_beta = []

for beta in betas:
    average_best_route_per_max_time = []

    for max_time in max_times:

        total_best_route_length = 0

        for _ in range(n_runs):
            print("Run number: ", _, "for max_time: ", max_time, " and beta: ", beta)
            clusters = split_tour_across_clusters(cities, 1)

            total_distance_for_run = 0

            for i, cluster_indices in enumerate(clusters.values()):
                cluster_cities = [cities[index] for index in cluster_indices]

                ant_colony = AntColony(cluster_cities, n_ants=n_ants, max_time=max_time, alpha=alpha, beta=beta, evaporation=evaporation, intensification=intensification)
                best_route = ant_colony.run()
                total_distance_for_run += total_distance(best_route)

            total_best_route_length += total_distance_for_run

        average_best_route_length = total_best_route_length / n_runs
        average_best_route_per_max_time.append(average_best_route_length)

    average_best_route_per_max_time_per_beta.append(average_best_route_per_max_time)

# Maintenant, nous avons toutes les valeurs moyennes, créons un histogramme
colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k']

total_bars_per_group = len(betas)
total_groups = len(max_times)

# Unité d'espace pour chaque groupe de barres (incluant l'espace entre les groupes)
unit_space_per_group = 1

# Espace qui serait occupé par les barres dans chaque groupe
bars_space_per_group = total_bars_per_group / (total_bars_per_group + 1)

# La largeur de chaque barre serait
bar_width = bars_space_per_group / total_bars_per_group
index = np.arange(len(max_times))

plt.figure()

for i, average_best_route_per_max_time in enumerate(average_best_route_per_max_time_per_beta):
    plt.bar(index + i * bar_width, average_best_route_per_max_time, bar_width, color=colors[i], label='Beta = '+str(betas[i]))
    for j, v in enumerate(average_best_route_per_max_time):
        plt.text(j + i * bar_width, v + 0.01 + (i*0.2), int(v), va='bottom', ha='center')

plt.axhline(y=optimal, color='r')
plt.xlabel('Max time')
plt.ylabel('Average best route length')

title = ""
title += "Comparaison of betas with average best route length ({} iterations) for different max times and each betas\n".format(n_runs)
title += "Nb cities: " + str(len(cities)) + " / "
title += "Ants: " + str(n_ants) + " / "
title += "Alpha: " + str(alpha) + " / "
title += "Evaporation: " + str(evaporation) + " / "
title += "Intensification: " + str(intensification)
plt.title(title)

plt.xticks(index + bar_width / 2, max_times)
plt.legend()

plt.tight_layout()
plt.show()