import math
import random

random.seed(42)
import itertools

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

# this class represents one solution in the search space
class TSPSolution:
    def __init__(self, points):
        self.points = points

    # compute the score of a tour
    def path_length(self):
        return sum(
            euclidean_distance(self.points[i], self.points[i + 1])
            for i in range(len(self.points) - 1)
        ) + euclidean_distance(self.points[-1], self.points[0])

    # choose one random neighbor by picking two random cities and reversing the entire
    # tour between these two cities
    def tweak(self):
        indices = sorted(random.sample(list(range(1, len(self.points))), 2))
        p = self.points
        new_points = (
            p[: indices[0]] + p[indices[0] : indices[1] + 1][::-1] + p[indices[1] + 1 :]
        )
        return TSPSolution(new_points)

    # computes the set of all neighbors (all possible tweaks)
    # this is not actually used in this code
    def neighborhood(self):
        new_solutions = []
        for indices in itertools.combinations(range(1, len(self.points)), 2):
            p = self.points
            new_points = (
                p[: indices[0]]
                + p[indices[0] : indices[1] + 1][::-1]
                + p[indices[1] + 1 :]
            )
            new_solutions.append(TSPSolution(new_points))
        return new_solutions
    
    # computes the set of all neighbors (all possible tweaks) and their scores
    # at the same time, using the fact that each tweak only removes two edges
    # and adds two edges
    def neighborhood_with_score(self):
        sol = self.points
        old_score = self.path_length()

        neighbors_with_scores = []
        for indices in itertools.combinations(range(1,len(sol)),2):
            score_change = 0
            p = sol
            new_points = p[:indices[0]] + p[indices[0]:indices[1]+1][::-1] + p[indices[1]+1:]

            score_change -= euclidean_distance(p[indices[0]-1],p[indices[0]])
            score_change -= euclidean_distance(p[indices[1]],p[(indices[1]+1) % len(p)])
            score_change += euclidean_distance(p[indices[0]-1],p[indices[1]])
            score_change += euclidean_distance(p[indices[0]],p[(indices[1]+1) % len(p)])
            
            neighbors_with_scores.append([TSPSolution(new_points), old_score + score_change])

        return neighbors_with_scores

# this class represents an instance of traveling saleman that we are trying to solve
class TSPProblem:
    def __init__(self, points):
        self.points = set(points)

    # generate a random solution by putting the cities in a random order
    def random_solution(self):
        return TSPSolution(random.sample(list(self.points), len(self.points)))

    # generate a random greedy solution by starting at a random city and then
    # always choosing the closest city as the next city to visit
    def random_greedy_solution(self):
        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 TSPSolution(solution)


def hill_climbing(tsp_problem, verbose=False):
    cur_sol = tsp_problem.random_solution()
    cur_score = cur_sol.path_length()

    failures = 0
    # we will stop when 10000 *consecutive* tweaks do not find a better solution
    while failures < 100_000:
        new_sol = cur_sol.tweak()
        new_score = new_sol.path_length()

        failures += 1

        if new_score < cur_score:
            cur_sol = new_sol
            cur_score = new_score
            if verbose:
                print(f"New best solution with score: {cur_score} (after {failures} failures)")
            failures = 0

    return cur_sol

def steepest_ascent_hill_climbing(tsp_problem, verbose=False):
    cur_sol = tsp_problem.random_solution()
    iter = 0

    while True:
        iter += 1
        neighbors_with_scores = cur_sol.neighborhood_with_score()
        best_neighbor, best_score = min(neighbors_with_scores, key=lambda x: x[1])
        if best_score >= cur_sol.path_length():
            # no improvement, we're done!
            return cur_sol
        if verbose:
            print(f"Iter {iter}: score = {best_score}")
        cur_sol = best_neighbor


def n_trials_steepest_ascent_hill_climbing(tsp_problem, n, verbose=False):
    cur_sol = tsp_problem.random_solution()
    iter = 0

    while True:
        iter += 1

        n_tweaks = [cur_sol.tweak() for _ in range(n)]
        best_tweak = min(n_tweaks, key=lambda x: x.path_length())
        if best_tweak.path_length() < cur_sol.path_length():
            cur_sol = best_tweak

        if verbose:
            print(f"Iter {iter}: score = {cur_sol.path_length()}")


def hill_climbing_with_random_restarts(tsp_problem):
    best_sol = None
    best_score = None

    iter = 0
    while True:
        iter += 1
        print(f"HHwRR iteration {iter}")
        sol = hill_climbing(tsp_problem)
        score = sol.path_length()

        if best_sol is None or score < best_score:
            best_sol = sol
            best_score = score

        print(f"\tResult of this iteration: {score}")
        print(f"\tBest solution found so far: {best_score}")

    return best_sol

# pick 300 random points in the unit square to form our input_data
points = [(random.random(), random.random()) for _ in range(300)]
tsp_problem = TSPProblem(points)

hill_climbing(tsp_problem, verbose=True)
# steepest_ascent_hill_climbing(tsp_problem, verbose=True)
# hill_climbing_with_random_restarts(tsp_problem)
# n_trials_steepest_ascent_hill_climbing(tsp_problem, 1000, verbose=True)