import json
import copy
from functools import reduce

import numpy as np  # contains helpful math functions like numpy.exp()
import numpy.random as random  # see numpy.random module
# import random  # alternative to numpy.random module

from typing import Tuple, List, Any

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
%matplotlib inline

"""Read input data and define helper functions for visualization."""

# Map services and data available from U.S. Geological Survey, National Geospatial Program.
# Please go to http://www.usgs.gov/visual-id/credit_usgs.html for further information
map = mpimg.imread('img/map.png')  # US States & Capitals map

# List of 30 US state capitals and corresponding coordinates on the map
with open('data/capitals.json', 'r') as capitals_file:
    capitals = json.load(capitals_file)
capitals_list = list(capitals.items())

def show_path(path, starting_city, w=12, h=8):
    """Plot a TSP path overlaid on a map of the US States & their capitals."""
    x, y = list(zip(*path))
    _, (x0, y0) = starting_city
    plt.imshow(map)
    plt.plot(x0, y0, 'y*', markersize=15)  # y* = yellow star for starting point
    plt.plot(x + x[:1], y + y[:1])  # include the starting point at the end of path
    plt.axis("off")
    fig = plt.gcf()
    fig.set_size_inches([w, h])

"""
Explanation: I. Introduction

As illustrated in the lectures, simulated annealing is a probablistic technique used for finding an approximate solution to an optimization problem.  In this exercise you will check your understanding by implementing simulated annealing to solve the Traveling Salesman Problem (TSP) between US state capitals.  Briefly, the TSP is an optimization problem that seeks to find the shortest path passing through every city exactly once.  In our example the TSP path is defined to start and end in the same city (so the path is a closed loop).
Image Source: Simulated Annealing - By Kingpin13 (Own work) [CC0], via Wikimedia Commons (Attribution not required)
Overview
Students should read through the code, then:

Implement the simulated_annealing() main loop function in Section II
Complete the TravelingSalesmanProblem class by implementing the successors() and get_value() methods in section III
Complete the schedule() function to define the temperature schedule in Section IV
Use the completed algorithm and problem description to experiment with simulated annealing to solve larger TSP instances on the map of US capitals
End of explanation
"""


def simulated_annealing(problem, schedule, rand=False):
    """The simulated annealing algorithm, a version of stochastic hill climbing
    where some downhill moves are allowed. Downhill moves are accepted readily
    early in the annealing schedule and then less often as time goes on. The
    schedule input determines the value of the temperature T as a function of
    time. [Norvig, AIMA Chapter 3]
    
    Parameters
    ----------
    problem : Problem
        An optimization problem, already initialized to a random starting state.
        The Problem class interface must implement a callable method
        "successors()" which returns states in the neighborhood of the current
        state, and a callable function "get_value()" which returns a fitness
        score for the state. (See the `TravelingSalesmanProblem` class below
        for details.)

    schedule : callable
        A function mapping time to "temperature". "Time" is equivalent in this
        case to the number of loop iterations.
    
    Returns
    -------
    Problem
        An approximate solution state of the optimization problem
        
    Notes
    -----
        (1) DO NOT include the MAKE-NODE line from the AIMA pseudocode

        (2) Modify the termination condition to return when the temperature
        falls below some reasonable minimum value (e.g., 1e-10) rather than
        testing for exact equality to zero
        
    See Also
    --------
    AIMA simulated_annealing() pseudocode
        https://github.com/aimacode/aima-pseudocode/blob/master/md/Simulated-Annealing.md
    """     
    
    """
    function SIMULATED-ANNEALING(problem,schedule) returns a solution state
     inputs: problem, a problem
        schedule, a mapping from time to "temperature"

     current ← MAKE-NODE(problem.INITIAL-STATE)
     for t = 1 to ∞ do
       T ← schedule(t)
       if T = 0 then return current
       next ← a randomly selected successor of current
       ΔE ← next.VALUE - current.VALUE
       if ΔE > 0 then current ← next
       else current ← next only with probability eΔE/T
       
    return cooler than 1e-10
    """
    time = 0
    current = problem
    while True:
        temp = schedule(time)
        time = time + 1
        if temp < 1e-10:
            return current
        if rand:
            # random switch successor is better
            successor = current.rand_successor()
        else:
            successor = random.choice(current.successors())
        delta = successor.get_value() - current.get_value()
        if delta > 0 or np.exp(delta / temp) > random.uniform(0.0, 1.0):
            current = successor

"""
Explanation: II. Simulated Annealing -- Main Loop
The main loop of simulated annealing repeatedly generates successors in the neighborhood of the current state and considers moving there according to an acceptance probability distribution parameterized by a cooling schedule.  See the simulated-annealing function pseudocode from the AIMA textbook online at github.  Note that our Problem class is already a "node", so the MAKE-NODE line is not required.
End of explanation
"""


Node = Tuple[str, Tuple[int, int]]

class TravelingSalesmanProblem:
    """Representation of a traveling salesman optimization problem.  The goal
    is to find the shortest path that visits every city in a closed loop path.
    
    Students should only need to implement or modify the successors() and
    get_values() methods.
    
    Parameters
    ----------
    cities : list
        A list of cities specified by a tuple containing the name and the x, y
        location of the city on a grid. e.g., ("Atlanta", (585.6, 376.8))
    
    Attributes
    ----------
    names
    coords
    path : list
        The current path between cities as specified by the order of the city
        tuples in the list.
    """
    def __init__(self, cities) -> None:
        self.path = copy.deepcopy(cities)
    
    def copy(self) -> Any:
        """Return a copy of the current board state."""
        new_tsp = TravelingSalesmanProblem(self.path)
        return new_tsp
    
    @property
    def names(self) -> Tuple:
        """Strip and return only the city name from each element of the
        path list. For example,
            [("Atlanta", (585.6, 376.8)), ...] -> ["Atlanta", ...]
        """
        names, _ = zip(*self.path)
        return names
    
    @property
    def coords(self) -> Tuple:
        """Strip the city name from each element of the path list and return
        a list of tuples containing only pairs of xy coordinates for the
        cities. For example,
            [("Atlanta", (585.6, 376.8)), ...] -> [(585.6, 376.8), ...]
        """
        _, coords = zip(*self.path)
        return coords
    
    def successors(self) -> List[Any]:
        """Return a list of states in the neighborhood of the current state by
        switching the order in which any adjacent pair of cities is visited.
        
        For example, if the current list of cities (i.e., the path) is [A, B, C, D]
        then the neighbors will include [A, B, D, C], [A, C, B, D], [B, A, C, D],
        and [D, B, C, A]. (The order of successors does not matter.)
        
        In general, a path of N cities will have N neighbors (note that path wraps
        around the end of the list between the first and last cities).

        Returns
        -------
        list<Problem>
            A list of TravelingSalesmanProblem instances initialized with their list
            of cities set to one of the neighboring permutations of cities in the
            present state
        """
        problems = []
        for i in range(len(self.path)):
            new_path = copy.deepcopy(self.path)
            j = i - 1
            if i == 0:
                j = len(new_path) - 1
            new_path[i], new_path[j] = new_path[j], new_path[i]
            problems.append(TravelingSalesmanProblem(new_path))
        return problems

    def rand_successor(self) -> List[Any]:
        new_path = copy.deepcopy(self.path)        
        i = random.randint(1, len(new_path) - 1)
        j = i
        while j == i:
            j = random.randint(1, len(new_path) - 1)
        new_path[i], new_path[j] = new_path[j], new_path[i]
        return TravelingSalesmanProblem(new_path)
    
    def get_value(self) -> float:
        """Calculate the total length of the closed-circuit path of the current
        state by summing the distance between every pair of adjacent cities.  Since
        the default simulated annealing algorithm seeks to maximize the objective
        function, return -1x the path length. (Multiplying by -1 makes the smallest
        path the smallest negative number, which is the maximum value.)
        
        Returns
        -------
        float
            A floating point value with the total cost of the path given by visiting
            the cities in the order according to the self.cities list
        
        Notes
        -----
            (1) Remember to include the edge from the last city back to the
            first city
            
            (2) Remember to multiply the path length by -1 so that simulated
            annealing finds the shortest path
        """
        distances = []
        coords = self.coords
        for i in range(len(coords)):
            j = i - 1
            if i == 0:
                j = len(coords) - 1
            # pythagoras theorem Euclidean distance
            a2 = (coords[i][0] - coords[j][0]) ** 2
            b2 = (coords[i][1] - coords[j][1]) ** 2
            c = np.sqrt(a2 + b2)
            distances.append(c)
        return sum(distances) * -1.

"""
Explanation: III. Representing the Problem
In order to use simulated annealing we need to build a representation of the problem domain.  The choice of representation can have a significant impact on the performance of simulated annealing and other optimization techniques.  Since the TSP deals with a close loop that visits each city in a list once, we will represent each city by a tuple containing the city name and its position specified by an (x,y) location on a grid.  The state will then consist of an ordered sequence (a list) of the cities; the path is defined as the sequence generated by traveling from each city in the list to the next in order.
End of explanation
"""


# Construct an instance of the TravelingSalesmanProblem
test_cities = [('DC', (11, 1)), ('SF', (0, 0)), ('PHX', (2, -3)), ('LA', (0, -4))]
test_names = ('DC', 'SF', 'PHX', 'LA')
test_coords = ((11, 1), (0, 0), (2, -3), (0, -4))
tsp = TravelingSalesmanProblem(test_cities)
assert(tsp.path == test_cities)
assert(tsp.names == test_names)
assert(tsp.coords == test_coords)

# Test the successors() method -- no output means the test passed
successor_paths = [x.path for x in tsp.successors()]
assert(all(x in [[('LA', (0, -4)), ('SF', (0, 0)), ('PHX', (2, -3)), ('DC', (11, 1))],
                 [('SF', (0, 0)), ('DC', (11, 1)), ('PHX', (2, -3)), ('LA', (0, -4))],
                 [('DC', (11, 1)), ('PHX', (2, -3)), ('SF', (0, 0)), ('LA', (0, -4))],
                 [('DC', (11, 1)), ('SF', (0, 0)), ('LA', (0, -4)), ('PHX', (2, -3))]]
          for x in successor_paths))

# Test the get_value() method -- no output means the test passed
assert(np.allclose(tsp.get_value(), -28.97, atol=1e-3))

"""
Explanation: Testing TravelingSalesmanProblem
The following tests should validate the class constructor and functionality of the successors() and get_value() methods.
End of explanation
"""


# These are presented as globals so that the signature of schedule()
# matches what is shown in the AIMA textbook; you could alternatively
# define them within the schedule function, use a closure to limit
# their scope, or define an object if you would prefer not to use
# global variables
alpha = 0.95
temperature=1e4

def schedule(time):
    # T(t) =α^t x T0
    return np.power(alpha, time) * temperature

"""
Explanation: IV. Define the Temperature Schedule
The most common temperature schedule is simple exponential decay:
$T(t) = \alpha^t T_0$
(Note that this is equivalent to the incremental form $T_{i+1} = \alpha T_i$, but implementing that form is slightly more complicated because you need to preserve state between calls.)
In most cases, the valid range for temperature $T_0$ can be very high (e.g., 1e8 or higher), and the decay parameter $\alpha$ should be close to, but less than 1.0 (e.g., 0.95 or 0.99).  Think about the ways these parameters effect the simulated annealing function.  Try experimenting with both parameters to see how it changes runtime and the quality of solutions.
You can also experiment with other schedule functions -- linear, quadratic, etc.  Think about the ways that changing the form of the temperature schedule changes the behavior and results of the simulated annealing function.
End of explanation
"""


# test the schedule() function -- no output means that the tests passed
assert(np.allclose(alpha, 0.95, atol=1e-3))
assert(np.allclose(schedule(0), temperature, atol=1e-3))
assert(np.allclose(schedule(10), 5987.3694, atol=1e-3))

# Failure implies that the initial path of the test case has been changed
assert(tsp.path == [('DC', (11, 1)), ('SF', (0, 0)), ('PHX', (2, -3)), ('LA', (0, -4))])
result = simulated_annealing(tsp, schedule)
print("Initial score: {}\nStarting Path: {!s}".format(tsp.get_value(), tsp.path))
print("Final score: {}\nFinal Path: {!s}".format(result.get_value(), result.path))
assert(tsp.path != result.path)
assert(result.get_value() > tsp.get_value())

"""
Explanation: Testing the Temperature Schedule
The following tests should validate the temperature schedule function and perform a simple test of the simulated annealing function to solve a small TSP test case
End of explanation
"""


# Create the problem instance and plot the initial state
num_cities = 30
capitals_tsp = TravelingSalesmanProblem(capitals_list[:num_cities])
starting_city = capitals_list[0]
print("Initial path value: {:.2f}".format(-capitals_tsp.get_value()))
print(capitals_list[:num_cities])  # The start/end point is indicated with a yellow star
show_path(capitals_tsp.coords, starting_city)

# set the decay rate and initial temperature parameters, then run simulated annealing to solve the TSP
alpha = 0.99
temperature=1e20
result = simulated_annealing(capitals_tsp, schedule, True)
print("Final path length: {:.2f}".format(-result.get_value()))
print(result.path)
show_path(result.coords, starting_city)

"""
Explanation: V. Run Simulated Annealing on a Larger TSP
Now we are ready to solve a TSP on a bigger problem instance by finding a shortest-path circuit through several of the US state capitals.
You can increase the num_cities parameter up to 30 to experiment with increasingly larger domains.  Try running the solver repeatedly -- how stable are the results?
End of explanation
"""


