# -*- coding: latin-1 -*-

from tspaco import Ant, World
from random import sample, uniform

inf = 1.0e34

class MinMaxAnt(Ant):
	def __init__(self, world, start):
		Ant.__init__(self, world, start)
		
	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 ** self.world.alpha) * (eta ** self.world.beta))
			except ZeroDivisionError:
				tcp.append(inf)
		return tcp
		
	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]
			
			
class MinMaxWorld(World):
	def __init__(self, distances, locations, ant_type=MinMaxAnt,
		num_ants=0, q_0=0.9, beta=2, rho=0.05, alpha=1.0, tau_0=0.0,
		global_best=False, cl_size=0, max_factor=5):
		
		World.__init__(self, distances, locations, ant_type, num_ants, q_0, beta, rho, alpha, tau_0, global_best, cl_size)
		
		self.tau_min = 2.0 / self.nn_heuristic()
		self.tau_max = max_factor * self.tau_min
		
		for row in xrange(len(self.pheromones)):
			for col in xrange(len(self.pheromones[row])):
				self.pheromones[row][col] = self.tau_max
		
	def __call__(self):
		self.min_max_step()
		
	def update_global_pheromones(self, tour, distance):
		r = tour[0]
		for s in tour[1:] + [r]:
			self.pheromones[r][s] = (1 - self.rho) * self.pheromones[r][s] + 1.0 / distance
			r = s
		
	def min_max_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)
		
		self.evaporate()

		if not self.locations_known:
			self.locations = self.compute_locations()

		if new_best:
			self.pretty_print()