""" Uses Simulated Annealing to find a (not necessarily optimal) solution for the Traveling Salesman Problem"""

import copy
import math
import matplotlib.pyplot as plt
import numpy as np
import random
import sys
import time
plt.ion()


def init_random(N):
	"""Generates two lists of random integers between 1 and 50"""
	X = [random.randint(1, 50) for i in range(N)]
	Y = [random.randint(1, 50) for i in range(N)]
	return (X,Y)


def tour_dist(X, Y):
	"""Computes the tour distance of N cities represented by their X and Y coordinates"""
	num_cities = len(X)
	dist = 0;
	for i in range(num_cities):
		start_x = X[i]
		start_y = Y[i]
        
		if i == num_cities-1:
			dest_x = X[0]
			dest_y = Y[0]
		else:
			dest_x = X[i+1]
			dest_y = Y[i+1]
		dist = dist + math.sqrt( math.pow((start_x-dest_x), 2) + math.pow((start_y-dest_y), 2))
	return dist

def simulated_annealing(curr_X, curr_Y):
	"""Implementation of Simulated Annealing for the Traveling Salesman Problem"""

	if len(curr_X) != len(curr_Y):
		print('X and Y must have same length')
		return

	N = len(curr_X)

	# The cooling schedule
	temp = 10;
	cooling_rate = .999;

	iterations = 0
	curr_cost = tour_dist(curr_X, curr_Y);
	
	while temp > .2:
		print("Temp=%.3f" % temp)		
		# Randomly select two cities to swap
		next_X = copy.copy(curr_X)
		next_Y = copy.copy(curr_Y)
		city1 = random.randint(0,N-1)
		city2 = random.randint(0,N-1)
        
		# Swap these two cities
		tmp = next_X[city1]
		next_X[city1] = next_X[city2]
		next_X[city2] = tmp
    
		tmp = next_Y[city1]
		next_Y[city1] = next_Y[city2]
		next_Y[city2] = tmp
    
		# Compute the cost of the new tour
		new_cost = tour_dist(next_X, next_Y)
    
		# New tour is better than curr tour. Automatically accept
		if (new_cost <= curr_cost): 
			curr_X = next_X
			curr_Y = next_Y            
			curr_cost = new_cost
            
		# New tour is worse...accept with some probability     
		elif random.random() <= math.exp((curr_cost - new_cost)/temp):			
			print('\tAccepted with prob %f' % math.exp((curr_cost - new_cost)/temp));
			curr_X = next_X
			curr_Y = next_Y
			curr_cost = new_cost

		# At end of iteration, update temp and plot
		temp = temp * cooling_rate
		iterations = iterations + 1

		if iterations % 100 == 0:
			plt.title('Iteration %d with cost %.3f' % (iterations, curr_cost))			
			plt.plot(curr_X, curr_Y, marker='x' )
			plt.plot([curr_X[-1], curr_X[0]], [curr_Y[-1], curr_Y[0]], marker='x',linestyle='--')
			plt.draw()
			plt.pause(0.8)
			plt.clf()
	
	# Plot the final tour
	plt.title('Final tour found with cost %.3f' % curr_cost)			
	plt.plot(curr_X, curr_Y, marker='x' )
	plt.plot([curr_X[-1], curr_X[0]], [curr_Y[-1], curr_Y[0]], marker='x',linestyle='--')
	plt.draw() 
	plt.pause(3) # pause so we can look at the solution
			

if __name__ == "__main__":
	print(sys.argv)
	if len(sys.argv) != 2:
		print('Usage: python simulated_annealing_demo.py [num_restarts]')
		sys.exit()
	
	num_cities = 20
	num_tries = int(sys.argv[1])

	# Randomly generate (x,y) coordinates for cities
	(x,y) = init_random(num_cities)

	for i in range(num_tries):
		simulated_annealing(x,y)