# AUTOGENERATED! DO NOT EDIT! File to edit: 07_tsp.ipynb (unless otherwise specified).

__all__ = ['dist', 'distanceGenerate', 'sortWaypoint', 'solve_tsp_nearest_neighbor', 'solve_tsp_held_karp']

# Cell

import sys
import itertools
import random
import time
import matplotlib.pyplot as plt
import numpy as np
import math
from python_tsp.heuristics import solve_tsp_simulated_annealing

def dist(p1, p2):
    """
        calculate the distance between two waypoints.
        Args:
            p1 : point's (x,y) position.
            p2 : point's (x,y) position.
        Returns:
            distance between two points.
    """
    return math.sqrt(((p1-p2)**2).sum())

def distanceGenerate(point_set):
    """
        generate a distance matrix based on the waypoints.
        Args:
            point_set : a set that contains all waypoint.
        Returns:
            a square matrix, which shows the distance between pairs of waypoint.
    """
    return np.asarray([[dist(np.array(p1), np.array(p2)) for p2 in point_set] for p1 in point_set])

def sortWaypoint(permutation, point_set):
    """
        accoriding to permutation calculated by tsp solver, return new set of waypoint which is sorted.
        Args:
            permutation : the order of waypoints calculated by tsp solver.
            point_set : a set that contains all waypoint.
        Returns:
            new set of waypoint which is sorted.
    """
    return [x for _, x in sorted(zip(permutation, point_set))]

def solve_tsp_nearest_neighbor(distance_matrix):
    """
        calculate tsp problem based on nearest neighbor, an algorithm that solves tsp using greedy assumption.
        Args:
            distance_matrix : a square matrix, which shows the distance between pairs of waypoint.
        Returns:
            A tuple, (path, cost)
            cost : optimal cost of tsp
            path : a orderd list of waypoint index based on distance matrix.
    """
    path = [0]
    cost = 0
    N = distance_matrix.shape[0]
    mask = np.ones(N, dtype=bool)
    mask[0] = False

    for i in range(N-1):
        last = path[-1]
        next_ind = np.argmin(distance_matrix[last][mask]) # find minimum of remaining locations
        next_loc = np.arange(N)[mask][next_ind] # convert to original location
        path.append(next_loc)
        mask[next_loc] = False
        cost += distance_matrix[last, next_loc]
        if(i == N-2):
            cost += distance_matrix[next_loc, 0]

    return path, cost

def solve_tsp_held_karp(distance_matrix):
    """
        calculate tsp problem based on Held-Karp, an algorithm that solves tsp using dynamic programming with memoization.
        Args:
            distance_matrix : a square matrix, which shows the distance between pairs of waypoint.
        Returns:
            A tuple, (path, cost)
            cost : optimal cost of tsp
            path : a orderd list of waypoint index based on distance matrix.
    """
    n = len(distance_matrix)
    C = {}

    # Set transition cost from initial state
    for k in range(1, n):
        C[(1 << k, k)] = (distance_matrix[0][k], 0)

    # Iterate subsets of increasing length and store intermediate results
    # in classic dynamic programming manner
    for subset_size in range(2, n):
        for subset in itertools.combinations(range(1, n), subset_size):
            bits = 0
            for bit in subset:
                bits |= 1 << bit
            # Find the lowest cost to get to this subset
            for k in subset:
                prev = bits & ~(1 << k)
                res = []
                for m in subset:
                    if m == 0 or m == k:
                        continue
                    res.append((C[(prev, m)][0] + distance_matrix[m][k], m))
                C[(bits, k)] = min(res)
    bits = (2**n - 1) - 1

    # Calculate optimal cost
    res = []
    for k in range(1, n):
        res.append((C[(bits, k)][0] + distance_matrix[k][0], k))
    opt, parent = min(res)

    # Backtrack to find full path
    path = []
    for i in range(n - 1):
        path.append(parent)
        new_bits = bits & ~(1 << parent)
        _, parent = C[(bits, parent)]
        bits = new_bits

    # Add implicit start state
    path.append(0)

    return list(reversed(path)), opt

