import matplotlib.pyplot as plt
import random, time
from libs.clustering import split_tour_across_clusters
from libs.simulated_annealing_stats import SimulatedAnnealing, total_distance

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

nb_ville = 100
max_coords = 1000
nb_truck = 1
temperature =  10000
cooling_rates = [ 0.999 , 0.99, 0.9 , 0.8]
temperature_ok = 0.001

start_time_generate = time.time()
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]]   
cities[0] = [max_coords/2, max_coords/2]
stop_time_generate = time.time()

optimal = 33523

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
    '#9467bd',  # Violet
    '#d62728',  # Rouge
    '#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
]


results = []

for cooling_rate in cooling_rates:
    print(f"\n---- Running with cooling rate: {cooling_rate} ----")
    
    distances_over_time = []
    
    for i, cluster_indices in enumerate(clusters.values()):
        cluster_cities = [cities[index] for index in cluster_indices]
        simulated_annealing = SimulatedAnnealing(cluster_cities, temperature=10000, cooling_rate=cooling_rate, temperature_ok=0.01)
        best_route, distances = simulated_annealing.run()
        distances_over_time.extend(distances)

    # Record results for this cooling rate
    results.append({
        'cooling_rate': cooling_rate,
        'distances': distances_over_time,
    })

# Plotting total distances for each cooling rate over time
plt.figure()
for result in results:
    plt.plot(result['distances'], label=f'Cooling rate: {result["cooling_rate"]}')
plt.xlabel('Iteration')
plt.ylabel('Total distance')
plt.legend(loc='upper right')
plt.title('Total distance over iterations for different cooling rates')
plt.axhline(y=optimal, color='black', linestyle='--')
plt.show()