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

name = "att48"
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

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

bar_width = 0.15
opacity = 0.8

x = np.arange(len(max_times))

fig, ax = plt.subplots()

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))
        ax.bar(x + (i * len(beta) + j) * bar_width, best_distances, bar_width, alpha=opacity, 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 / 2)
ax.set_xticklabels(max_times)
ax.legend()

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