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

name = "berlin52"
cities = [[565, 575], [25, 185], [345, 750], [945, 685], [845, 655], [880, 660], [25, 230], [525, 1000], [580, 1175], [650, 1130], [1605, 620], [1220, 580], [1465, 200], [1530, 5], [845, 680], [725, 370], [145, 665], [415, 635], [510, 875], [560, 365], [300, 465], [520, 585], [480, 415], 
[835, 625], [975, 580], [1215, 245], [1320, 315], [1250, 400], [660, 180], [410, 250], [420, 555], [575, 665], [1150, 1160], [700, 580], [685, 595], [685, 610], [770, 610], [795, 645], [720, 635], [760, 650], [475, 960], [95, 260], [875, 920], [700, 500], [555, 815], [830, 485], 
[1170, 65], [830, 610], [605, 625], [595, 360], [1340, 725], [1740, 245]]
optimal = 7542

n_ants = 10
alpha = [1, 2]
beta = [2, 3, 4, 5]
evaporation = 0.5
intensification = 2
max_times = [1]
iterations = 1

bar_width = 0.15
opacity = 0.8

x = np.arange(len(max_times))

fig, ax = plt.subplots()

# Preparing the colormap
colors = [
    '#1f77b4',  # Bleu moyen
    '#ff7f0e',  # Orange
    '#2ca02c',  # Vert
    '#d62728',  # Rouge
    '#9467bd',  # Violet
    '#8c564b',  # Marron
    '#e377c2',  # Rose
    '#7f7f7f',  # Gris
    '#bcbd22',  # Vert olive
    '#17becf',  # Turquoise
    '#1b9e77',  # Vert Teal
    '#d95f02',  # Orange foncé
    '#7570b3',  # Violet moyen
    '#e7298a',  # Fuchsia
    '#66a61e',  # Vert pomme
    '#e6ab02',  # Jaune or
    '#a6761d',  # Bronze
    '#666666',  # Gris foncé
    '#f781bf',  # Rose clair
    '#999999',  # Gris moyen
]
color_dict = {}

for i in range(len(alpha)):
    for j in range(len(beta)):
        best_distances = []
        for max_time in max_times:
            for iteration in range(iterations):
                ant_colony = AntColony(cities, n_ants, alpha[i], beta[j], evaporation, intensification, max_time)
                print("Running iteration number {}/{} ({} sec)".format(iteration + 1, iterations, max_time))
                best_route = ant_colony.run()
                best_distances.append(total_distance(best_route))
        
        # Defining a unique key for each alpha-beta pair
        key = f"alpha={alpha[i]}_beta={beta[j]}"
        if key not in color_dict:
            color_dict[key] = colors[i * len(beta) + j]
        
        ax.bar(x + (i * len(beta) + j) * bar_width, best_distances, bar_width, alpha=opacity, color=color_dict[key], label='alpha={} beta={}'.format(alpha[i], beta[j]))

ax.set_xlabel('Max Time')
ax.set_ylabel('Best Distance')
ax.set_title('Best distances per max time for each alpha and beta')
ax.set_xticks(x + bar_width / 3)
ax.set_xticklabels(max_times)
ax.legend()

plt.axhline(y=optimal, color='r')
plt.show()