import psp2d, pspos
 
WHITE_COLOR = psp2d.Color(255,255,255)
CLEAR_COLOR = psp2d.Color(0,0,0,255)
RED_COLOR   = psp2d.Color(255, 0, 0)
 
cm     = []
coords = []
CITIES = 20
 
pspos.setclocks(333,166)
psp_scr  = psp2d.Screen()
psp_font = psp2d.Font('font.png')
psp_scr.clear(CLEAR_COLOR)
psp_font.drawText(psp_scr, 0, 5, "Loading Pyevolve modules...")
psp_scr.swap()
 
from pyevolve import G1DList
from pyevolve import GSimpleGA
from pyevolve import GAllele
from pyevolve import Mutators
from pyevolve import Crossovers
from pyevolve import Consts
 
from random import shuffle as rand_shuffle, randint as rand_randint
from math import sqrt
 
def cartesian_matrix(coords):
   """ A distance matrix """
   matrix={}
   for i,(x1,y1) in enumerate(coords):
      for j,(x2,y2) in enumerate(coords):
         dx, dy = x1-x2, y1-y2
         dist=sqrt(dx*dx + dy*dy)
         matrix[i,j] = dist
   return matrix
 
def tour_length(matrix, tour):
   """ Returns the total length of the tour """
   total = 0
   for i in range(CITIES):
      j      = (i+1)%CITIES
      total += matrix[tour[i], tour[j]]
   return total
 
def write_tour_to_img(coords, tour):
   """ The function to plot the graph """
   psp_scr.clear(CLEAR_COLOR)
   for i in range(CITIES):
      j      = (i+1)%CITIES
      city_i = tour[i]
      city_j = tour[j]
      x1, y1 = coords[city_i]
      x2, y2 = coords[city_j]
      psp_scr.drawLine(int(x1), int(y1), int(x2), int(y2), WHITE_COLOR)
      psp_font.drawText(psp_scr, int(x1)+7, int(y1)-5, str(i))
      psp_scr.fillRect(int(x1), int(y1), 6, 6, RED_COLOR)
 
   psp_scr.swap()
 
def G1DListTSPInitializator(genome, **args):
   """ The initializator for the TSP """
   lst = [i for i in xrange(genome.getListSize())]
   rand_shuffle(lst)
   genome.genomeList = lst
 
def evolve_callback(ga_engine):
   """ Callback called every generation by Pyevolve """
   write_tour_to_img(coords, ga_engine.bestIndividual())
   return False
 
def main_run():
   global cm, coords
 
   width, height = psp_scr.size
   coords = [(rand_randint(0, width),rand_randint(0, height))
                 for i in xrange(CITIES)]
   cm     = cartesian_matrix(coords)
 
   setOfAlleles = GAllele.GAlleles(homogeneous=True)
   range_allele = GAllele.GAlleleRange(0, len(coords)-1)
   setOfAlleles.add(range_allele)
 
   genome = G1DList.G1DList(len(coords))
   genome.setParams(allele=setOfAlleles)
 
   genome.evaluator.set(lambda chromosome: tour_length(cm, chromosome))
   genome.mutator.set(Mutators.G1DListMutatorSwap)
   genome.crossover.set(Crossovers.G1DListCrossoverOX)
   genome.initializator.set(G1DListTSPInitializator)
 
   ga = GSimpleGA.GSimpleGA(genome)
   ga.setMinimax(Consts.minimaxType["minimize"])
   ga.setGenerations(300)
   ga.setPopulationSize(200)
   ga.setCrossoverRate(1.0)
   ga.setMutationRate(0.1)
   ga.stepCallback.set(evolve_callback)
 
   ga.evolve()
 
if __name__ == "__main__":
   main_run()