import sys
import copy
import math
from heapq import heappush, heappop
from random import randint
import matplotlib.pyplot as plt

distances = []
n = 0

# returns hueristic value
def heuristic(source, target, dist):
    #return 0
    for i in range(0, min(n * 2, n ** 2)):
        if (dist <= distances[i]):
            return 0
    return dist


# returns weight(distance) between 2 vertices
def weight(city1, city2):
    return math.sqrt((city1[1] - city2[1]) ** 2 + (city1[2] - city2[2]) ** 2)


# returns if going to city i from curnode is okay
def valid(curnode, cities, i):
    # can't revisit a city
    if (i in curnode[2]):
        return False
    # if curnode hasn't visited everything else and we are going to the gaol
    elif (len(curnode[2]) != (len(cities) - 1) and i == (len(cities) - 1)):
        return False
    else:
        return True


# returns the count of child states. updates frontier
def explore(curnode, frontier, cities):
    children = 0
    
    for i in range(1, len(cities)):
        # can we go here?
        if (valid(curnode, cities, i)):
            children += 1
            # add to frontier
            dist = weight(curnode[1], cities[i])
            metric = dist + heuristic(curnode[1], cities[i], dist)
            heappush(frontier, [curnode[0] + metric, cities[i], copy.deepcopy(curnode[2])])
    
    return children


# A* solver for TSP
def solve(frontier, cities):
    # number of child states
    newnodes = 0
    curnode = heappop(frontier)
    
    # while we are not on the goal node
    while (curnode[1] != cities[-1]):
        newnodes += explore(curnode, frontier, cities)
        curnode = heappop(frontier)
        curnode[2].append(cities.index(curnode[1]))
        
    return newnodes + 1


def tsp(filename):
    # import map
    #prob = "instance_10.txt"
    prob = filename
    with open(prob) as file:
        prob = file.read().splitlines()

    n = int(prob[0])
    cities = prob[1:]
    counter = 0
    # convert to list of list of ints
    for l in cities:
        temp = l.split()
        temp[1] = int(temp[1])
        temp[2] = int(temp[2])
        temp.append(counter)
        counter += 1
        cities[cities.index(l)] = temp
        
    distances = []
    
    for i in range(0, len(cities) - 1):
        for j in range(i, len(cities)):
            dist = weight(cities[i], cities[j])
            distances.append(dist)
                
    distances.sort()
    
    # add in goal state
    cities.append(copy.deepcopy(cities[0]))
    cities[-1][0] = "Goal"
    
    # create frontier
    frontier = []
    # add start state, list of visited, and distancesofar to frontier
    heappush(frontier, [0, cities[0], [0]])

    result = solve(frontier, cities)
    print(result)
    return result


def main():
    global distances
    global n
    plt.ioff()
    plt.switch_backend('agg')
    averages = []
    
    for i in range(1, 12):
        average = 0
        for j in range(1, 11):
            filepath = "tsp_problems/" + str(i) + "/instance_" + str(j) + ".txt"
            average += tsp(filepath)
            
        averages.append(average / 10.0)
    
    figure, axes = plt.subplots(1, 1, True)
    axes.semilogy(range(1, 12), averages, label='TSP Solver (Heuristic)')
    axes.legend()
    plt.xlabel("Number of Cities")
    plt.ylabel("Average Number of Nodes Generated in 10 Runs")
    plt.savefig("tsph.pdf")
    
main()