from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import numpy as np
import random, time
from clustering import split_tour_across_clusters

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

def plot_clusters(cities, clusters):
    # Création d'une liste de couleurs pour les différents clusters
    colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k']

    # Création d'un nouveau graphique
    plt.figure()

    # Pour chaque cluster
    for i, cluster in clusters.items():
        # Sélection d'une couleur pour le cluster
        color = colors[i % len(colors)]

        # Pour chaque ville dans le cluster
        for city_index in cluster:
            # Récupération des coordonnées de la ville
            city = cities[city_index]

            # Ajout de la ville au graphique
            plt.scatter(city[0], city[1], c=color, s=20)

    # show first city in black and twice bigger
    plt.scatter(cities[0][0], cities[0][1], c='k', s=200)


    # Affichage du graphique
    plt.show()


nb_ville = 100
max_coords = 1000
nb_truck = 5

# Define the coordinates of the cities
# And set depot at the first city in the middle of the map
start_time_generate = time.time()
cities = generate_cities(nb_ville, max_coords)
cities[0] = [max_coords/2, max_coords/2]
stop_time_generate = time.time()

# Split the tour across clusters with nb_truck trucks
start_time_split = time.time()
clusters = split_tour_across_clusters(cities, nb_truck)
stop_time_split = time.time()

# show the number of cities in each cluster
for cluster in clusters.values():
    print(len(cluster))

# show the time
print("\n---- TIME ----")
print("generate cities time: ", stop_time_generate - start_time_generate)
print("split cities time: ", stop_time_split - start_time_split)

# show the clusters
plot_clusters(cities, clusters)