# -*- coding: latin-1 -*-
#
# Ant Colony System for Symmetric and Asymmetric TSPs
# By: Geoffrey Foster
# School of Computer Science, Carleton University
# © 2007, Geoffrey Foster
#
###############################################################################
#
# Based upon "Solving Symmetric and Asymmetric TSPs by Ant Colonies" by
# Luca Maria Gambardella and Marco Dorigo
# International Conference on Evolutionary Computation, 1996
# pp. 622-627
# http://citeseer.ist.psu.edu/gambardella96solving.html
#
#
###############################################################################
# η(r,s) is a heuristic function which evaluates the utility of move s when in city r. In
# ATSP problems it is the inverse of the distance from r to s
#
# τ(r,s) is a positive real value associated with the edge (r,s) and is the measure of
# pheromones left on a trail
#
# ß weighs the relative importance of the heuristic function
#
# An agent (ant) situated in city r moves to city s using the following pseudo-random-proportional
# action choice rule (or state transition rule).
#     ⎧arg_{u ∈ J_k(r)} max {[τ(r,s)]⋅[η(r,s)]^ß} if q ≤ q_0 (exploration)
# s = ⎨
#     ⎩S			otherwise (exploitation)
#
# q is chosen at random, uniformly between [0,1], q_0 is a paramater (0 ≤ q_0 ≤ 1)
#
#
# S is a random variable selected according to the distribution given by the following formula
# which gives the probability with which an agent in city r chooses the city s to move to.
#                ⎧ [τ(r,s)⋅(η(r,s)^ß)]/sum_{z ε J_k(r)}[τ(r,z)]⋅[η(r,z)]^ß	if s ε J_k(r)
# p_k(r,s) = S = ⎨
#                ⎩ 0		otherwise
#
# J_k(r) is the list of cities still to be visited
#
# ρ is the evaporation co-efficient 0 ≤ ρ ≤ 1
#
###############################################################################
#
# Initial defaults:
# n = number of cities
# q_0 = 0.9, ß = 2, ρ = α = 0.1, m = 10,
# and τ_0 = a very small constant (n⋅L_{nn})^{-1} where L_{nn} is produced by the nearest neighbor heuristic
# cl_size is the maximum size of the candidate list for each vertex
#
###############################################################################

from math import cos, sin, pi
from operator import itemgetter
from random import choice, sample, uniform
from tsplib.atsp import full_matrix

inf = 1.0e34

class Ant(object):
	def __init__(self, world, start):
		self.world = world
		self.start = start
		self.reset()
	
	def reset(self):
		self.tour = [self.start]
		self.remaining = filter(lambda x: x not in self.tour, range(len(self.world.locations)))
		self.tour_distance = 0
		
	def __repr__(self):
		return '<Ant: Distance=%s; Tour=%s>' % (self.tour_distance, self.tour)
		
	def __call__(self):
		current = self.tour[-1]
		if len(self.remaining):
			cl = filter(lambda x: x not in self.tour, self.world.candidate_list[current])
			q = uniform(0.0, 1.0)
			if len(cl) > 0:
				next = cl[0]
			elif q <= self.world.q_0:
				next = self.explore()
			else:
				next = self.exploit()
			
			self.tour_distance += self.world.distances[current][next]
			
			self.remaining.remove(next)
			self.tour.append(next)
		else:
			self.tour_distance += self.world.distances[current][self.start]

	def explore(self):
		"""
		arg-max_{u ∈ J_k(r)} {[τ(r,s)]⋅[η(r,s)]^ß}
		"""
		current = self.tour[-1]
		
		tcp = self.trail_closeness_product()
		
		# Find the maximum trail-closeness product
		index = None
		for i in xrange(len(tcp)):
			if index is None or tcp[i] > tcp[index]:
				index = i
		
		return self.remaining[index]
				
	def trail_closeness_product(self):
		"""
		Compute trail–closeness product for each s ∈ J_k(r)
		"""
		current = self.tour[-1]
		tcp = []
		for city in self.remaining:
			tau = self.world.pheromones[current][city]
			try:
				eta = (1.0 / self.world.distances[current][city])
				tcp.append(tau * (eta ** self.world.beta))
			except ZeroDivisionError:
				#eta = inf
				tcp.append(inf)
				#tcp.append(float('inf'))
		return tcp
		
	def exploit(self):
		"""
		S is a random variable selected according to the distribution given by the following formula
		which gives the probability with which an agent in city r chooses the city s to move to.
		               ⎧ [τ(r,s)⋅(η(r,s)^ß)]/sum_{z ε J_k(r)}[τ(r,z)]⋅[η(r,z)]^ß	if s ε J_k(r)
		p_k(r,s) = S = ⎨
		               ⎩ 0		otherwise
		
		J_k(r) is the list of cities still to be visited
		"""
		current = self.tour[-1]
		
		# Compute numerator of p_k(r,s) for each s ∈ J_k(r)
		tcp = self.trail_closeness_product()
		
		# Compute the denominator
		tcp_sum = sum(tcp)
		
		# Now create a probability distribution list
		distribution = []
		total = 0.0
		for p in tcp:
			total += p/tcp_sum
			distribution.append(total)
		
		# Pick a uniform random number between [0.0, 1.0) and walk the list
		# looking for correct p_k(r,s). Once found return it.
		u = uniform(0.0, 1.0)
		index = 0
		for i in xrange(len(distribution)):
			index = i
			if u <= distribution[i]:
				break
		
		return self.remaining[index]
		

class World(object):
	def __init__(self, distances, locations, ant_type=Ant,
		num_ants=0, q_0=0.9, beta=2, rho=0.1, alpha=0.1, tau_0=0.0, global_best=True, cl_size=0):
		
		self.q_0, self.beta, self.rho, self.alpha, self.tau_0 = q_0, beta, rho, alpha, tau_0
		self.global_best = global_best
		self.cl_size = cl_size
		
		self.distances = distances

		if tau_0 == 0.0:
			#tau_0 = 1.0 / (len(self.distances) * self.nn_heuristic())
			self.tau_0 = 10**-6
			#tau_0 = 1.0
			
		# Initialize the global pheromones between edges to nothing
		pheromones = []
		for row in xrange(len(self.distances)):
			pheromones.append([])
			for col in xrange(len(self.distances[row])):
				pheromones[row].append(self.tau_0)
		self.pheromones = pheromones
		
		# Setup the (x,y) coordinates for the locations
		# If we don't know them then flag as unknown
		self.locations_known = True
		for location in locations.values():
			x,y = location
			if x is None or y is None:
				self.locations_known = False
				break
		self.locations = locations
		
		if not num_ants:
			num_ants = len(self.locations)
		
		start_points = []
		while num_ants > len(self.locations):
			start_points.extend(xrange(len(self.locations)))
			num_ants -= len(self.locations)
			
		start_points.extend(sample(xrange(len(self.locations)), num_ants))
		self.ants = [ant_type(self, point) for point in start_points]
		
		self.best_solution = (None, None)
		
		self.candidate_list = self.build_cl(cl_size)
		
	def build_cl(self, size):
		cl = []
		for i in xrange(len(self.distances)):
			icl = [(index, self.distances[i][index]) for index in xrange(len(self.distances))]
			icl.sort(key=itemgetter(1))
			
			cl.append([value[0] for value in icl[:size]])
		return cl
			
		
	def size(self):
		return len(self.locations)
		
	def __repr__(self):
		return repr(self.ants) + '\n' + repr(self.pheromones)
		
	def nn_heuristic(self):
		distance = 0.0
		vertices = range(len(self.distances))
		o = s = choice(vertices)
		vertices.remove(s)
		while len(vertices) > 0:
			r = None
			for v in vertices:
				if r is None or self.distances[s][v] < self.distances[s][r]:
					r = v
			vertices.remove(r)
			distance += self.distances[s][r]
			s = r
		distance += self.distances[s][o]
		return distance
		
	def evaporate(self):
		for r in xrange(len(self.pheromones)):
			for s in xrange(len(self.pheromones)):
				self.pheromones[r][s] *= (1.0 - self.rho)
				self.ensure_minimum(r, s)

	def ensure_minimum(self, r, s):
		if self.pheromones[r][s] < self.tau_0:
			self.pheromones[r][s] = self.tau_0

	def update_local_pheromones(self):
		"""
		Does both updating and decaying on a tours pheromones
		"""
		for ant in self.ants:
			r,s = ant.tour[-2:]
			self.pheromones[r][s] = (1.0 - self.rho) * self.pheromones[r][s] + self.rho * self.tau_0
			self.ensure_minimum(r, s)
		
	def update_global_pheromones(self, tour, distance):
		"""
		Does both updating and decaying of the global pheromones
		"""
		r = tour[0]
		for s in tour[1:] + [r]:
			self.pheromones[r][s] = (1.0 - self.alpha) * self.pheromones[r][s] + (self.alpha / distance)
			self.ensure_minimum(r, s)
			r = s

	def add_local_pheromones(self):
		"""
		Does just the addition of pheromones to local tour
		"""
		for ant in self.ants:
			r,s = ant.tour[-2:]
			self.pheromones[r][s] += self.tau_0

	def add_global_pheromones(self, tour, distance):
		"""
		Does just the addition of pheromones to the global tour
		"""
		r = tour[0]
		for s in tour[1:] + [r]:
			self.pheromones[r][s] += (self.alpha / distance)
			r = s
			
	def __call__(self):
		#self.dorigo_step()
		self.dorigo_step_modified()
		
	def dorigo_step(self):
		new_best = False
		for ant in self.ants:
			ant.reset()
			
		for i in xrange(len(self.locations) + 1):
			# Have each ant compute its trail
			for ant in self.ants:
				ant()
			
			# Update local pheromones
			self.update_local_pheromones()	
		
		# Store the best tour
		best = self.best()
		if best[1] != self.best_solution[1]:
			new_best = True
			self.best_solution = self.best()
		btour, bdistance = self.best_solution
		
		# Could update based upon global best or iteration best
		# if iteration best then L_{best-iter} is length of best iteration
		# if global best then L_{best-iter} is length of best overall
		if self.global_best:
			# Update the global pheromones based on best overall
			tour = btour
			distance = bdistance
		else:
			# Update the global pheromones based on best of iteration
			itour, idistance = self.iteration_best()
			tour = itour
			distance = idistance
			
		self.update_global_pheromones(tour, distance)
		
		if not self.locations_known:
			self.locations = self.compute_locations()
			
		if new_best:
			self.pretty_print()
		
	def dorigo_step_modified(self):
		"""
		Same as the original Dorigo algorithm but instead of combining both updating of
		pheromones and decay in one step the local and global pheromones are applied first
		and then at the end a global decay occurs
		"""
		new_best = False
		for ant in self.ants:
			ant.reset()
			
		for i in xrange(len(self.locations) + 1):
			# Have each ant compute its trail
			for ant in self.ants:
				ant()
			
			# Update local pheromones
			self.add_local_pheromones()
		
		# Store the best tour
		best = self.best()
		if best[1] != self.best_solution[1]:
			new_best = True
			self.best_solution = self.best()
		btour, bdistance = self.best_solution
		
		# Could update based upon global best or iteration best
		# if iteration best then L_{best-iter} is length of best iteration
		# if global best then L_{best-iter} is length of best overall
		if self.global_best:
			# Update the global pheromones based on best overall
			tour = btour
			distance = bdistance
		else:
			# Update the global pheromones based on best of iteration
			itour, idistance = self.iteration_best()
			tour = itour
			distance = idistance
			
		self.add_global_pheromones(tour, distance)
		
		self.evaporate()
		
		if not self.locations_known:
			self.locations = self.compute_locations()
			
		if new_best:
			self.pretty_print()
			
	def iteration_best(self):
		shortest = None
		for ant in self.ants:
			if shortest is None or ant.tour_distance < shortest.tour_distance:
				shortest = ant
		
		return (shortest.tour[:], shortest.tour_distance)
		
	def best(self):
		iter_best = self.iteration_best()
		itour, idistance = iter_best
		
		if self.best_solution == (None, None):
			return iter_best
		else:
			btour, bdistance = self.best_solution
			if idistance < bdistance:
				return iter_best
			else:
				return self.best_solution
		
		
	def compute_locations(self):
		"""
		If the locations of each city are already known then return them,
		otherwise:
			Compute locations for each city based upon a circle layout.
			Cities are spaced out according to distance
			Locations are shifted so that the point (0,0) is top-left
		"""
		if self.locations_known:
			return self.locations
			
		tour, distance = self.best_solution
		C = distance # circumference
		radius = C / (2 * pi)
		arc = 0.0
		locations = {tour[0]: (2 * radius, radius)}
		index = 1
		r = tour[0]
		for s in tour[1:]:
			arc += self.distances[r][s]
			theta = arc / radius # arc length divided by radius = angle
			x = cos(theta) * radius
			y = sin(theta) * radius
			
			locations[s] = (radius + x, radius - y)
			index += 1
			r = s
		
		return locations
		
	def pretty_print(self):
		tour, distance = self.best_solution
		if tour is None or distance is None:
			return
		
		s = tour[0]
		s_tour = '%s' % s
		for r in tour[1:] + [s]:
			s_tour += ' -> %s' % r
			s = r
		print 'Total Distance: %s' % distance
		print s_tour
		
def main():
	f=open('./data/ALL_atsp/br17.atsp')
	#f=open('./data/ALL_atsp/ft53.atsp')
	matrix, coords = full_matrix(f)
	world = World(distances=matrix, locations=coords)
	#print world.build_cl(10)
	print world.ants[0].trail_closeness_product()
	world()
	
	#print world.pheromones
	#for i in xrange(world.size() * 1000):
	#	world()
	#	print world.best()
	#print world.pheromones
		
	
if __name__ == '__main__':
	main()