from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import numpy as np
import random, time, math
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 distance(city1, city2):
    return math.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2)

def total_distance(cities):
    return sum([distance(cities[i - 1], cities[i]) for i in range(len(cities))])

previous_route = None

def simulated_annealing(cities, temperature=10000, cooling_rate=0.9999, temperature_ok=0.001, cluster_index=0):
    interration = 0
    current_solution = cities.copy()
    best_solution = cities.copy()
    while temperature > temperature_ok:
        new_solution = current_solution.copy()
        # Swap two cities in the route
        i = random.randint(0, len(new_solution) - 1)
        j = random.randint(0, len(new_solution) - 1)
        new_solution[i], new_solution[j] = new_solution[j], new_solution[i]
        # Calculate the acceptance probability
        current_energy = total_distance(current_solution)
        new_energy = total_distance(new_solution)
        delta = new_energy - current_energy
        if delta < 0 or random.random() < math.exp(-delta / temperature):
            current_solution = new_solution
            if total_distance(current_solution) < total_distance(best_solution):
                best_solution = current_solution
        # Cool down
        temperature *= cooling_rate
        interration += 1
        # Print every 1000 iterations
        if interration % 1000 == 0:
            print("Cluster", cluster_index, ": iteration", interration, "with current total distance", total_distance(current_solution))
    return best_solution


nb_ville = 20
max_coords = 1000
nb_truck = 4
temperature = 10000
cooling_rate = 0.999
temperature_ok = 0.001

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()

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

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 simulated_annealing
    best_route = simulated_annealing(cluster_cities, temperature, cooling_rate, temperature_ok)
    best_routes.append((best_route, color))

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

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=color, marker='x', linestyle='-', label=f"Cluster {i}")
plt.legend(loc="best")
plt.show()