#! /usr/bin/env python3
#
def five_test ( ):

  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'five_test():' )
  print ( '  TSP for a set of five cities.' )

  position = np.loadtxt ( 'five_position.txt' )
  city_num = position.shape[0]
  names = [ 'Aspen', 'Boston', 'Clayton', 'Dayton', 'Easton' ]
#
#  Plot the cities.
#
  from graphviz import Graph
  dot = Graph ( format = 'png' )
  for n in names:
    dot.node ( n )
  for i in range ( 0, city_num ):
    for j in range ( i + 1, city_num ):
      dot.edge ( names[i], names[j] )
  dot.render ( 'five.dot', view = True )
  filename = 'five.dot.png'
  print ( '  Graphics saved as "%s"' % ( filename ) )
#
#  Compute distance matrix.
#
  distance = np.zeros ( [ city_num, city_num ] )
  for i in range ( 0, city_num ):
    for j in range ( 0, city_num ):
      distance[i,j] = np.round ( np.linalg.norm ( position[i,:] - position[j,:] ) )

  print ( '' )
  print ( '  Distance table:' )
  print ( distance )
#
#  Plot the cities with distances.
#
  from graphviz import Graph
  dot = Graph ( format = 'png' )
  for n in names:
    dot.node ( n )
  for i in range ( 0, city_num ):
    for j in range ( i + 1, city_num ):
      dot.edge ( names[i], names[j], str ( distance[i,j] ) )
  dot.render ( 'five_distance.dot', view = True )
  filename = 'five_distance.dot.png'
  print ( '  Graphics saved as "%s"' % ( filename ) )
#
#  Brute force solution.
#
  from itertools import permutations
 
  perms = permutations ( np.arange ( city_num ) )

  mileage_shortest = np.inf

  for perm in perms:
    mileage = path_length ( distance, perm )
    if ( mileage < mileage_shortest ):
      perm_shortest = perm
      mileage_shortest = mileage

  print ( '' )
  print ( '  Shortest path is ', perm_shortest )
  print ( '  Shortest mileage is ', mileage_shortest )

  for i in range ( 0, city_num ):
    j = perm_shortest[i]
    print ( names[j], end = '' )
    print ( ' => ', end = '' )
  j = perm_shortest[0]
  print ( names[0] )

  return

def path_length ( distance, perm ):
  import numpy as np
  mileage = 0.0
  c0 = perm[-1]
  for c1 in perm:
    mileage = mileage + distance[c0,c1]
    c0 = c1
  return mileage

five_test ( )