from random import Random
from time import time
import math
import inspyred

class tsp:
  
  def __init__ ( self, weights ):
    # weights contains the distances between all cities
    self.weights = weights
    # TSP is a minimization problem
    self.maximize = False
    # use a predefined bounder function, in which only values in range(len(weights)) are allowed
    self.bounder = inspyred.ec.DiscreteBounder([i for i in range(len(weights))])

  # a generator generates one individual
  def generator(self,random, args):
    locations = [i for i in range(len(self.weights))]
    random.shuffle(locations)
    return locations
  
  # an evaluator returns a list of fitnesses for all individuals in a population
  def evaluator(self, candidates, args):
    fitness = []
    for candidate in candidates:
       total = 0
       for src, dst in zip(candidate, candidate[1:] + [candidate[0]]):
         total += self.weights[src][dst]
       fitness.append(total)
    return fitness


def main(prng=None, display=False):    
    if prng is None:
        prng = Random()
        prng.seed(time()) 
        
    # small example of cities
    points = [(110.0, 225.0), (161.0, 280.0), (325.0, 554.0), (490.0, 285.0), 
              (157.0, 443.0), (283.0, 379.0), (397.0, 566.0), (306.0, 360.0), 
              (343.0, 110.0), (552.0, 199.0)]
    # calculate distances
    weights = [[0 for _ in range(len(points))] for _ in range(len(points))]
    for i, p in enumerate(points):
        for j, q in enumerate(points):
            weights[i][j] = math.sqrt((p[0] - q[0])**2 + (p[1] - q[1])**2)
              
    # initialize the problem
    problem = tsp ( weights )
    ea = inspyred.ec.EvolutionaryComputation(prng)
    
    # make a choice for selection, in this case tournament selection
    ea.selector = inspyred.ec.selectors.tournament_selection
    
    # make a selection for the operations, in this case PMX crossover and
    # inversion mutation. See the online documentation for options
    ea.variator = [inspyred.ec.variators.partially_matched_crossover, 
                   inspyred.ec.variators.inversion_mutation]
                   
    # the offspring obtained after applying the above operators in order,
    # replace the parents in the population; however, the best individual
    # in the parent population will also survive if it is better than the
    # worst individual in the new population (due to num_elites)
    ea.replacer = inspyred.ec.replacers.generational_replacement
    
    # stopping criterion
    ea.terminator = inspyred.ec.terminators.generation_termination
    
    # call of the evolutionary algorithm; note that the parameters
    # are passed in a dictionary. This dictionary will be forwarded
    # to all of the above functions (args), which can use information
    # passed here, such as num_elites for the elitism used in 
    # ea.replacer.
    final_pop = ea.evolve(generator=problem.generator, 
                          evaluator=problem.evaluator, 
                          bounder=problem.bounder,
                          maximize=problem.maximize, 
                          pop_size=100, 
                          max_generations=50,
                          tournament_size=5,
                          num_selected=100,
                          num_elites=1)
    
    if display:
        best = max(ea.population)
        print('Best Solution: {0}: {1}'.format(str(best.candidate), best.fitness))
    return ea
            
if __name__ == '__main__':
    main(display=True)