from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import numpy as np
import random, time, math
from libs.clustering import split_tour_across_clusters
from libs.aco import AntColony, total_distance

def generate_cities(nb, max_coords=1000):
    return [random.sample(range(max_coords), 2) for _ in range(nb)]

nb_ville = 1500
max_coords = 1000
nb_truck = 2
max_time = 10
nb_ants = 10

max_time_per_cluster = max_time / nb_truck

start_time_generate = time.time()
cities = generate_cities(nb_ville, max_coords)
stop_time_generate = time.time()


start_time_split = time.time()
clusters = split_tour_across_clusters(cities, nb_truck)
stop_time_split = time.time()

def elbow_method(cities, max_clusters):
    inertias = []
    for i in range(1, max_clusters+1):
        kmeans = KMeans(n_clusters=i, random_state=0).fit(cities)
        inertias.append(kmeans.inertia_)
    return inertias

cities = generate_cities(nb_ville, max_coords) # Assurez-vous que les villes sont générées avant de les passer à la méthode du coude
inertias = elbow_method(cities, max_clusters=50)

plt.figure()
plt.plot(range(1, len(inertias)+1), inertias, marker='o')
plt.title('Elbow Method')
plt.xlabel('Number of clusters')
plt.ylabel('Inertia')
plt.show()




for cluster in clusters.values():
    print(len(cluster))
print("\n---- TIME ----")
print("generate cities time: ", stop_time_generate - start_time_generate)
print("split cities time: ", stop_time_split - start_time_split)

# create new figure for annealing paths
plt.figure()
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
]

best_routes = []

for i, cluster_indices in enumerate(clusters.values()):
    # Sélection d'une couleur pour le cluster
    color = colors[i % len(colors)]
    
    # Récupération des coordonnées de la ville
    cluster_cities = [cities[index] for index in cluster_indices]

    # Appel de la fonction AntColony.run
    ant_colony = AntColony(cluster_cities, n_ants=nb_ants, max_time=max_time_per_cluster, alpha=1, beta=5)
    best_route = ant_colony.run()
    best_routes.append((best_route, color))

    print("Total distance for cluster", i, ": ", total_distance(best_route))

# calculate total distance for all clusters
full_total_distance = 0
for route, color in best_routes:
    full_total_distance += total_distance(route)

print("Total distance for all clusters: ", full_total_distance)

for i, (route, color) in enumerate(best_routes):
    x = [city[0] for city in route]
    y = [city[1] for city in route]
    x.append(x[0])
    y.append(y[0])
    plt.plot(x, y, color="blue", marker='o', linestyle='-', label=f"Cluster {i}")
    # add title with nb_ville, nb_truck and max_time
    plt.title(f"nb_ville = {len(cities)}, nb_truck = {nb_truck}, max_time = {max_time}")   

plt.show()
