from sklearn.cluster import KMeans
import numpy as np

def split_tour_across_clusters(cities, nb_truck):
    if nb_truck == 1:
        return {0: list(range(len(cities)))}

    # clustering initial
    kmeans = KMeans(n_clusters=nb_truck, random_state=0).fit(cities)
    clusters = {i:[] for i in range(nb_truck)}

    # assignation des indices des villes aux clusters
    for i, label in enumerate(kmeans.labels_):
        clusters[label].append(i)

    max_iterations = len(cities)**2
    iteration = 0

    while True:
        iteration += 1
        if iteration > max_iterations:
            print("Le nombre maximum d'itérations a été atteint. La boucle a été interrompue.")
            break
        # calcul des tailles de clusters
        cluster_sizes = {i:len(clusters[i]) for i in range(nb_truck)}

        # identification du cluster le plus grand et du plus petit
        max_cluster = max(cluster_sizes, key=cluster_sizes.get)
        min_cluster = min(cluster_sizes, key=cluster_sizes.get)

        # s'il n'y a pas de grande disparité, on arrête la boucle
        if cluster_sizes[max_cluster] - cluster_sizes[min_cluster] <= 1:
            break

        # calcul du centre de chaque cluster
        cluster_centers = {i:np.mean([cities[index] for index in clusters[i]], axis=0) for i in range(nb_truck)}

        # calcul des distances entre le centre du cluster le plus grand et les autres
        distances = {i:np.linalg.norm(cluster_centers[max_cluster]-cluster_centers[i]) for i in range(nb_truck)}
        del distances[max_cluster]  # on supprime la distance vers lui-même

        if nb_truck >= 3:
            # on identifie les 2 clusters les plus proches
            closest_clusters = sorted(distances, key=distances.get)[:2]

            
            # parmi les deux clusters les plus proches, on choisit le plus petit
            if cluster_sizes[closest_clusters[0]] <= cluster_sizes[closest_clusters[1]]:
                target_cluster = closest_clusters[0]
            else:
                target_cluster = closest_clusters[1]
        else:
            closest_clusters = sorted(distances, key=distances.get)[:1]
            target_cluster = closest_clusters[0]



        # si le transfert va créer une plus grande disparité, on arrête la boucle
        if cluster_sizes[target_cluster] >= cluster_sizes[max_cluster]:
            break

        # calcul des distances entre le centre du cluster cible et les villes du cluster le plus grand
        distances_to_target = {index:np.linalg.norm(cluster_centers[target_cluster]-cities[index]) 
                               for index in clusters[max_cluster]}

        # on identifie la ville la plus proche du centre du cluster cible
        closest_city_index = min(distances_to_target, key=distances_to_target.get)

        # on transfère la ville du cluster le plus grand au cluster cible
        clusters[target_cluster].append(closest_city_index)
        clusters[max_cluster].remove(closest_city_index)

    # Ajout du point de départ et d'arrivée pour chaque cluster
    depot_index = 0
    for cluster in clusters.values():
        if cluster[0] != depot_index:
            cluster.insert(0, depot_index)
        if cluster[-1] != depot_index:
            cluster.append(depot_index)

    return clusters