# Tabu search for Traveling Salesman, 200 cities
# this version scores a new tour by just subtracting
# the two disappearing edges and adding the two new 
# edges

import math
import random
random.seed(1)
import itertools

from collections import defaultdict

def euclidean_distance(pt1, pt2):
    return math.sqrt((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2)

class TSP:
    """
    This is a class that represents a traveling salesman problem
    (not a solution, a whole problem)
    """

    def __init__(self, points):
        self.points = set(points)

    def random_solution(self):
        """
        returns a random solution (the point of the TSP in a random order)
        """
        return random.sample(list(self.points), len(self.points))

    def random_greedy_solution(self):
        """
        returns a greedy solution starting at a random point
        """
        start_point = random.choice(list(self.points))
        solution = [start_point]

        points_left = set(self.points)
        points_left.remove(start_point)

        while len(points_left) > 0:
            closest = min(
                points_left, key=lambda p: euclidean_distance(p, solution[-1])
            )
            solution.append(closest)
            points_left.remove(closest)

        return solution

    def score(self, sol):
        """
        scores a tour, from scratch
        """
        return sum(
            euclidean_distance(sol[i], sol[i + 1])
            for i in range(len(sol) - 1)
        ) + euclidean_distance(sol[-1], sol[0])

    def neighborhood_with_score(self, sol, old_score):
        """
        computes every neighbor of "sol" formed by picking two cities
        and reverse the whole block of the tour in between them

        for each new neighbor, it also computes what its score is and returns that
        """
        for indices in itertools.combinations(range(1,len(sol)),2):
            score_change = 0
            
            new_points = sol[:indices[0]] + sol[indices[0]:indices[1]+1][::-1] + sol[indices[1]+1:]

            score_change -= euclidean_distance(sol[indices[0]-1], sol[indices[0]])
            score_change -= euclidean_distance(sol[indices[1]], sol[(indices[1]+1) % len(sol)])
            score_change += euclidean_distance(sol[indices[0]-1], sol[indices[1]])
            score_change += euclidean_distance(sol[indices[0]], sol[(indices[1]+1) % len(sol)])
            
            yield new_points, old_score + score_change, indices

# form a random TSP problem
points = [(random.random(), random.random()) for _ in range(200)]
tsp = TSP(points)

# track best sol and best score
best_sol = None
best_score = None

# taboo_time = 100 means we won't allow swap the block of cities between i and j
#   if we've already done that in the last 100 moves
taboo_time = 100
# a "defaultdict" is a really use special kind of dictionary that you give a
#   type, in this case "int", and if you try to access a key that doesn't exist
#   then it will default to 0
taboo = defaultdict(int)

# random start solution
sol = tsp.random_solution()

generation = 0

# repeat until you get bored
while True:
    generation += 1

    # get score of current solution, save if best
    length = tsp.score(sol)
    if best_score is None or best_score > length:
        best_score = length
        best_sol = sol
    if generation > 1:
        print(f"\t\tBest ever score: {best_score}")
    
    # set up variables
    steepest_best_sol = None
    steepest_best_score = None
    steepest_move = None

    # for every possible move (every sol in the neighborhood)
    for s, pl, move in tsp.neighborhood_with_score(sol, length):
        # if taboo[move] > generation, then this move is still taboo and
        #   we can't do it. BUT, if it's taboo and its score is still better
        #   than anything we've ever seen before, go anyway
        if move in taboo and taboo[move] > generation and not pl < best_score:
            continue

        # if we can do it, then get its score, and see if it's the best
        #   thing in the neighborhood
        if steepest_best_score is None or pl < steepest_best_score:
            steepest_best_sol = s
            steepest_best_score = pl
            steepest_move = move

    # move to the best thing in the neighborhood, and make that move
    #   taboo for a while
    sol = steepest_best_sol
    print(f"Gen {generation} -- Cur score: {steepest_best_score}", end="")
    taboo[steepest_move] = generation + taboo_time
    

