import numpy as np
import matplotlib.pyplot as plt

# Question 1:
cities = np.array(['Paris', 'Lyon', 'Toulouse', 'Nice', 'Nantes', 'Montpellier', 'Strasbourg', 'Bordeaux', 'Lille', 'Le Havre', 'Clermont-Ferrand', 'Limoges', 'Orleans'])
longitudes = np.array([2. + 21./60. + 7./3600., 4. + 49./60. + 56./3600., 1. + 26./60. + 38./3600., 7. + 16./60. + 17./3600., - (1. + 33./60. + 10./3600.), 3. + 52./60. + 38./3600., 7. + 45./60. + 8./3600., - (0. + 34./60. + 46./3600.), 3. + 3./60. + 48./3600., 0. + 6./60., 3. + 4./60. + 56/3600., 1. + 15./60., 1. + 54./60. + 32./3600.])
latitudes = np.array([48. + 51./60. + 24./3600., 45. + 45./60. + 28./3600., 43. + 36./60. + 16./3600., 43. + 41./60. + 45/3600., 47. + 13./60. + 5./3600., 43. + 36./60. + 43./3600., 48. + 34./60. + 24./3600., 44. + 50./60. + 16./3600., 50. + 38./60. + 14./3600., 49. + 29./60. + 24./3600., 45. + 46./60. + 33./3600., 45. + 51./60., 47. + 54./60. + 9./3600.])

longitudes *= np.pi / 180. # Conversion in radians
latitudes *= np.pi / 180.
R = 6371 # Earth radius in km
distances = R * np.arccos(np.cos(latitudes) * np.cos(latitudes.reshape((cities.size, 1))) * np.cos(longitudes - longitudes.reshape(cities.size, 1)) + np.sin(latitudes) * np.sin(latitudes.reshape((cities.size, 1))))

# Question 2: the solution uses a loop, but you can avoid a loop by using the function roll of NumPy.
def path_distance(path, distances):
    total_length = 0.
    n = path.size
    for i in range(n):
        total_length += distances[path[i], path[(i + 1) % n]]
    return total_length
    

def path_distance_without_loop(path, distances):
    total_length = np.sum(distances[path, np.roll(path, -1)])
    return total_length
    
    
optimal_path = np.array([0, 12, 11, 10, 1, 3, 5, 2, 7, 4, 9, 8, 6])
Lmin = path_distance_without_loop(optimal_path, distances)
print(Lmin, path_distance(optimal_path, distances))

# Question 3:
def simulated_annealing(distances, alpha, Ti, Tf):
    n = np.shape(distances)[0]
    path = np.arange(n)
    while Ti > Tf:
        i = np.random.randint(1, n) # we propose to exchange the cities at location i and location j in the path (except Paris -> we only choose integers larger than 1).
        j = np.random.randint(1, n)
        while i == j: # we pick up an integer until i =! j.
            j = np.random.randint(1, n)
        if np.abs(i - j) > 1: # the calculation of the change in length of the path is different whether the two cities follow each other on the path.
            diff_length = - (distances[path[i], path[(i + 1) % n]] + distances[path[i], path[i - 1]] + distances[path[j], path[(j + 1) % n]] + distances[path[j], path[j - 1]]) + (distances[path[j], path[(i + 1) % n]] + distances[path[j], path[i - 1]] + distances[path[i], path[(j + 1) % n]] + distances[path[i], path[j - 1]]) # if (i,j) are switched, the length of the circuit is affected via the distance between cities i, j and their neighbours in the path.
        elif j == i + 1:
            diff_length = - (distances[path[j], path[(j + 1) % n]] + distances[path[i], path[i - 1]]) + (distances[path[i], path[(j + 1) % n]] + distances[path[j], path[i - 1]]) # j is the right neighbour of i, so only the distances with the left neighbour of i and the right neighbour of j are affected. 
        else:
            diff_length = - (distances[path[j], path[j - 1]] + distances[path[i], path[(i + 1) % n]]) + (distances[path[i], path[j - 1]] + distances[path[j], path[(i + 1) % n]]) # j is the left neighbour of i, so only the distances with the right neighbour of i and the left neighbour of j are affected. 
        if diff_length < 0: # Metropolis
            path[i], path[j] = path[j], path[i] # we exchange the two cities in the path.
        elif np.random.random() < np.exp(- diff_length / Ti):
            path[i], path[j] = path[j], path[i]
        Ti *= (1. - alpha) # we reduce the temperature.
    return path

# Question 4:
Ti = 100.
Tf = 1e-5
alpha_arr = np.logspace(-1, -4, 4)
length_sa_arr = np.zeros_like(alpha_arr)
nrepet = 20
for a, alpha in enumerate(alpha_arr):
    for k in range(nrepet):
        print(alpha, k)
        path_sa = simulated_annealing(distances, alpha, Ti, Tf)
        length_sa_arr[a] += path_distance_without_loop(path_sa, distances)
    length_sa_arr[a] /= nrepet
print(length_sa_arr)

# Question 5: the length of the path found is always larger than the length of the optimal path.
# So the optimal path is not found, but paths of length comparable are easily found in few seconds.
plt.figure()
plt.semilogx(alpha_arr, length_sa_arr, '-o')
plt.semilogx(alpha_arr, Lmin * np.ones_like(alpha_arr), '--')
plt.xlabel(r'$\alpha$')
plt.ylabel(r'$\langle L\rangle(\alpha)$')
plt.grid()
plt.show()

# Question 6: the precision scales as alpha^{1/3}.
plt.figure()
plt.loglog(alpha_arr, length_sa_arr - Lmin, '-o')
plt.loglog(alpha_arr, 1e3 * alpha_arr ** (1. / 3.), '--')
plt.xlabel(r'$\alpha$')
plt.ylabel(r'$\langle L\rangle(\alpha)-L_\mathrm{min}$')
plt.grid()
plt.show()