#!/usr/bin/python
# Copyright Joel Friedman 2018.  For use in UBC's Math 441 course.
#
# improvements over jf_tsp5.txt: (1) add user for the value of n,
#   (2) translate u variables by 1, (3) add output option (with ugly results...)

from gurobipy import *

######################################################## 
# Set up the variables and constraints
######################################################## 

n = input("How many cities? ")            # n is the number of cities

jfTSP = Model('Baby_TSP')

x = jfTSP.addVars(n,n, vtype=GRB.BINARY, name='ArcsTaken')
u = jfTSP.addVars(n, lb = 1, ub = n, vtype=GRB.INTEGER, name='order')

jfTSP.addConstrs( ( x[i,i] == 0   for i in range(n) ), name='ZeroDiag' )
jfTSP.addConstrs( ( x.sum(i,'*') == 1  for i in range(n) ), name='LeaveCities')
jfTSP.addConstrs( ( x.sum('*',i) == 1  for i in range(1,n) ), name='EnterCities')

jfTSP.addConstrs( ( u[i] - u[j] + n * x[i,j] <= n-1   
  for i in range(1,n)  for j in range(n) ),  name='orderRespected')

######################################################## 
# Optimize, display results
######################################################## 

jfTSP.setObjective( x[0,1] , GRB.MAXIMIZE )
jfTSP.optimize()
jfTSP.printAttr("X")

######################################################## 
# Append results to a file
######################################################## 

output_file_name = input("Name of file for appending (in quotation marks): ")
if ( output_file_name != "" ) :
    output_file = open( output_file_name , "a")     # open a file for appending
    for v in jfTSP.getVars():
        output_file.write(v.varName + " " + str( v.x ) )