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

# globals to contain reference solutions
neos14 = [318, 324, 336, 319, 351, 311, 272, 361, 274, 322]
neos15 = [313, 318, 281, 324, 378, 291, 348, 342, 353, 325]
neos16 = [404, 353, 361, 349, 358, 344, 373, 355, 243, 330]
neos36 = 464
neos = {14: neos14, 15: neos15, 16: neos16, 36: neos36}

# returns a shuffled version of the cities list while maintaining the first & last elements
def rand_order(cities):
    # shitty variable name
    suburbs = cities[1:-1]
    shuffle(suburbs)
    suburbs.insert(0, cities[0])
    suburbs.append(cities[-1])
    return suburbs


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


# returns total cost/distance of a tour (list of cities)
def length(cities):
    sum = 0
    for i in range(len(cities) - 1):
        sum += dist(cities[i], cities[i+1])

    return sum


# returns the best ordered child state of a state (hill climbing, basic) & its cost
def best_child(cities, cost):
    bestorder = cities
    bestcost = cost
    
    for i in range(1, len(cities) - 2):
        for j in range(i, len(cities) - 1):
            temporder = copy.deepcopy(cities)
            tempcity = temporder[i]
            temporder[i] = temporder[j]
            temporder[j] = tempcity
            templen = length(temporder)
            if (templen < cost):
                bestorder = temporder
                bestcost = templen
    
    return (bestorder, bestcost)


# returns the best child using tabu
# returns the best ordered child state of a state (hill climbing, basic) & its cost
def bester_child(cities, cost, tabu):
    bestorder = cities
    bestcost = cost
    
    for i in range(1, len(cities) - 2):
        for j in range(i, len(cities) - 1):
            temporder = copy.deepcopy(cities)
            tempcity = temporder[i]
            temporder[i] = temporder[j]
            temporder[j] = tempcity
            
            if (get_hashable(temporder) in tabu):
                continue
            
            templen = length(temporder)
            if (templen < cost):
                bestorder = temporder
                bestcost = templen
    
    return (bestorder, bestcost)


# returns a random child for simulated annealing search
# returns the ordering and the cost of it
def random_child(cities):
    #i = randbelow(len(cities) - 3) + 1
    i = randint(1, len(cities) - 2)
    j = i
    
    while (j == i):
        #j = randbelow(len(cities) - 3) + 1
        j = randint(1, len(cities) - 2)
        
    temporder = copy.deepcopy(cities)
    tempcity = temporder[i]
    temporder[i] = temporder[j]
    temporder[j] = tempcity
    
    return(temporder, length(temporder))


# returns the basic hill climbing solution for a given input
def sir_edmund_hillary(cities):
    curcost = length(cities)
    curpath = copy.deepcopy(cities)
    steps = 0
    
    # keep trying things til we can no longer try things
    while(1):
        steps += 1
        child = best_child(curpath, curcost)
        if (child[1] < curcost):
            curcost = child[1]
            curpath = child[0]
            continue
        break
        
    return(curpath, curcost, steps)


# retuns a tuple of the indices (ordered) in a list of cities
def get_hashable(cities):
    newlist = []
    for i in range(1, len(cities) - 1):
        newlist.append(cities[i][1])
        
    return tuple(newlist)


# returns the hill climbing solution w/tabu & sideways moves for a given input
def tenzing_norgay(cities):
    curcost = length(cities)
    curpath = copy.deepcopy(cities)
    tabu = {get_hashable(cities)}
    laterals = 0
    
    # keep trying things til we can no longer try things
    while(1):
        child = bester_child(curpath, curcost, tabu)
        if (child[1] < curcost):
            curcost = child[1]
            curpath = child[0]
            tabu.add(get_hashable(curpath))
            continue
        if (child[1] == curcost and laterals < n):
            laterals += 1
            curcost = child[1]
            curpath = child[0]
            tabu.add(get_hashable(curpath))
            continue
        break
        
    return(curpath, curcost)


# returns the hill climbing solution w/ x random restarts for a given input
# if you want x restarts, call with x + 1
def reinhold_messner(cities, x):
    curcost = length(cities)
    curpath = copy.deepcopy(cities)
    bestcost = length(cities)
    bestpath = copy.deepcopy(cities)
    
    # keep trying things til we can no longer try things
    while(x):
        child = best_child(curpath, curcost)
        if (child[1] < curcost):
            curcost = child[1]
            curpath = child[0]
            continue
        else:
            if (curcost < bestcost):
                bestcost = curcost
                bestpath = copy.deepcopy(curpath)
            curpath = rand_order(cities)
            curcost = length(curpath)
            x -= 1
        
    return(bestpath, bestcost)


# returns true or false randomly, based on the distribution e^(delta / temp)
def bad_weather(delta, temp):
    return random.SystemRandom().random() < math.exp(delta / temp)


# decrements temp by the option specified by opt: 1 is linear, 2 is logarithmic, 3 is exponential
def colder(temp, opt):
    if (opt == 1):
        return temp - 25
    if (opt == 2):
        return temp - (math.log(temp) + 2)
    else:
        return temp - math.exp(-0.001 * temp)


# returns the hill climbing solution using simulated annealing w/ schedule opt
def beck_weathers(cities, opt):
    temp = 1000
    curcost = length(cities)
    curpath = copy.deepcopy(cities)
    
    # keep trying things til we can no longer try things
    while(temp > 0):
        child = random_child(curpath)
        delta = curcost - child[1]
        if (delta > 0 or bad_weather(delta, temp)):
            curcost = child[1]
            curpath = child[0]
        temp = colder(temp, opt)
        
    return(curpath, curcost)


# solves the inputted TSP problem (filepath), using algorithm algo
# 1 = hill climbing, 2 = tabu/sideways allowed, 3 = random restarts, 4 = annealing
# x represents the optional parameters for 3 and 4
# specifically, the # of restarts and the annealing schedule, respectively
def solver(problem, algo, *x):
    # import map
    prob = problem
    with open(prob) as file:
        prob = file.read().splitlines()

    global n
    cities = prob[1:]
    counter = 0
    # convert to list of list of ints
    # original format is: letter, coord, coord
    # output format is: iD#, letter, coord, coord
    for l in cities:
        temp = l.split()
        temp[1] = int(temp[1])
        temp[2] = int(temp[2])
        temp.insert(0, counter)
        counter += 1
        cities[cities.index(l)] = temp
    
    # add in goal state
    cities.append(cities[0])
    input = rand_order(cities)
    
    if (algo == 1):
        result = sir_edmund_hillary(cities)
    elif (algo == 2):
        result = tenzing_norgay(cities)
    elif (algo == 3):
        result = reinhold_messner(cities, x)
    else:
        result = beck_weathers(cities, x)
    
    return result
    

# main function to get aggregate data and graph stuff
def main():
    plt.ioff()
    plt.switch_backend('agg')
    solution_steps = []
    solution_scores = []
    good_sol_counts = []
    
    for i in range(14, 17):
        steps = []
        solution_steps.append(steps)
        scores = []
        solution_scores.append(scores)
        good_sols = []
        good_sol_counts.append(good_sols)
        
        for j in range(1, 11):
            filepath = "tsp_problems/" + str(i) + "/instance_" + str(j) + ".txt"
            prob_steps = 0
            prob_scores = 0
            good_solutions = 0
            # run problem 100 times
            for k in range(0, 100):
                # path, cost, steps returned
                result = solver(filepath, 1)
                prob_steps += result[2]
                prob_scores += (result[1] / neos[i][j - 1])
                if (result[1] <= neos[i][j- 1]):
                    good_solutions += 1
            
            steps.append(prob_steps / 100.0)
            scores.append(prob_scores / 100.0)
            good_sols.append(good_solutions)
            
    figure, axes = plt.subplots(3, 2, True)
            
    for i in range(0, 3):
        axes[i][0].plot(range(1, 11), solution_steps[i], label = "Steps to Solution", color = 'b')
        axes[i][0].set_xlabel("Problem Instance w/ " + str(i + 14) + " cities")
        axes[i][0].set_ylabel("Steps Taken", color = 'b')
        axis_twin = axes[i].twinx()
        axis_twin.plot(range(1, 11), solution_scores[i], label = "Solution Quality", color = 'r')
        axis_twin.set_ylabel("Solution Quality", color = 'r')
        #axes[i].legend()
        #axis_twin.legend()
    for i in range(0, 3):
        axes[i][1].plot(range(1, 11), good_sol_counts[i], color = 'b')
        axes[i][1].set_xlabel("Problem Instance w/ " + str(i + 14) + " cities")
        axes[i][1].set_ylabel("% of runs <= NEOS", color = 'b')

    plt.savefig("hill_climbing.pdf")

main()