# Imports
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import LineCollection

"""
Explanation: The Trapped Knight sequence
Inspired by This Numberphile video.
Suppose you have an infinitely large chessboard. Number the squares in a spiral fashion (as you would in an Ulam spiral) and place a regular chess Knight in square 1. At every step you move the knight to the lowest-valued square that it hasn't visited yet. What will happen?
Relevant OEIS sequence: A316667--Squares visited by knight moves on a spirally numbered board and moving to the lowest available unvisited square at each step.
End of explanation
"""


# Moving in the cardinal directions:
# E: (+1, +0)
# S: (+0, -1)
# W: (-1, +0)
# N: (+0, +1)

def move_e():
    exes.append(exes[-1]+1)
    eyes.append(eyes[-1])

def move_s():
    exes.append(exes[-1])
    eyes.append(eyes[-1]-1)

def move_w():
    exes.append(exes[-1]-1)
    eyes.append(eyes[-1])

def move_n():
    exes.append(exes[-1])
    eyes.append(eyes[-1]+1)

"""
Explanation: The main strategy is as follows:

Build an odd-sized square lattice (grid),
Assign a value to every lattice point, starting from 1 and following a spiral. Keep track of every value's xand y coordinates in a pair of lists (exes, eyes),
Place a virtual knight on the lattice's middle point,
Keep track of the knight's movements in a pair of lists (path_x, path_y) that will contain respectively its x and y coordinates
Keep track of visited squares in another list (visited)
At every step, we will calculate:
The possible landing squares for the knight, given its current position, and
Which of those have not been visited, and 
The lowest of all legal, non-visited squares
Move the knight to the next square given by the rules.
???
PROFIT

Definitions
move_e(), move_s(), move_w, move_n()
These four mini functions will serve in building the number spiral and will help to keep track of the coordinates of every number in the sequence. The movements are named such that they parallel the cardinal directions of a map where the North points upward (since not all maps follow this convention). Thus, East is "to the right", West is "to the left", North is "up" and South is "Down".
End of explanation
"""


def move_where(step):
    '''
    defines where to move based on some mod 4 rule
    '''
    if step % 4 == 1:
        return move_e()
    elif step % 4 == 2:
        return move_s()
    elif step % 4 == 3:
        return move_w()
    else:
        return move_n()

"""
Explanation: move_where(step)
From sequence A063826 in the OEIS. Let the numbers 1, 2, 3 and 4 represent "steps" in the East, South, West and North directions respectively. How these numbers are derived will be seen later on.
End of explanation
"""


def move_k(pos, grid):
    '''
    Returns the 8 possible chess knight movements for specified position 
    in grid.
    
    'pos' must be a tuple, 'grid' is a specified 2-dimensional (numpy)
    array with labeled squares.
    '''
    moves = []
    k1 = grid[pos[0]+1, pos[1]+2]
    k2 = grid[pos[0]+2, pos[1]+1]
    k3 = grid[pos[0]+2, pos[1]-1]
    k4 = grid[pos[0]+1, pos[1]-2]
    k5 = grid[pos[0]-1, pos[1]-2]
    k6 = grid[pos[0]-2, pos[1]-1]
    k7 = grid[pos[0]-2, pos[1]+1]
    k8 = grid[pos[0]-1, pos[1]+2]
    moves.append(k1)
    moves.append(k2)
    moves.append(k3)
    moves.append(k4)
    moves.append(k5)
    moves.append(k6)
    moves.append(k7)
    moves.append(k8)
    return moves

"""
Explanation: move_k(pos, grid)
Given a position tuple in a particular grid, this will return the 8 possible moves a knight can make, as a series of tuples.
End of explanation
"""


def get_next(pos, grid):
    global visited
    # Step 1: Get all possible moves for post
    candidates1 = np.array(move_k(pos, grid))
    # Step 2: Return moves that are not in visited
    candidates2 = np.setdiff1d(candidates1, visited)
    # Step 3: Determine the lowest of those
    # TODO: Set up what happens if it's an empty array (20191020: fixed?)
    if candidates2.size > 0:
        winner = int(min(candidates2))
        # Step 4: Append this winner to the number of visited squares
        visited.append(winner)
        # Step 5: Determine coordinates of last winner
        next_x = exes[winner - 1]
        next_y = eyes[winner - 1]
    else:
        print('Sorry')  # For debugging purposes
        next_x = -1
        next_y = -1
    return next_x, next_y

"""
Explanation: get_next(pos, grid)
This function will look at the visited squares in the Knight path so far and will determine where to move next via:

Get all possible knight-moves starting from pos
Eliminate those that have been visited
Determine the lowest of those, assuming there are valid candidates left (if not, see step 5)
"Visit" the square, adding it to visited and determining its $(x, y)$ coordinates 
If there are no more valid candidates, pseudo-break the cycle
return the coordinates to the Knight's path
End of explanation
"""


def plot_moves_from_number(number):
    """
    Plots a single fracture point and all knight-moves from it
    """
    x = exes[number-1]
    y = eyes[number-1]
    xs = [x+1, x+2, x+2, x+1, x-1, x-2, x-2, x-1]
    ys = [y+2, y+1, y-1, y-2, y-2, y-1, y+1, y+2]
    # plt.scatter(x, y, marker='o', color='k')
    plt.text(x, y, str(number))
    plt.scatter(xs, ys, marker='^', color='r')

def plot_moves_from_position(pos):
    """
    Plots a single fracture point and all knight-moves from it
    """
    x = pos[0]
    y = pos[1]
    xs = [x+1, x+2, x+2, x+1, x-1, x-2, x-2, x-1]
    ys = [y+2, y+1, y-1, y-2, y-2, y-1, y+1, y+2]
    plt.scatter(x, y, marker='o', color='k')
    plt.scatter(xs, ys, marker='^', color='r')

"""
Explanation: Visualization functions
These functions will plot all valid knight moves from a given number in the Ulam spiral, or from a given position in the grid, assuming they are called in a valid matplotlib Figure.
End of explanation
"""


start = 1

"""
Explanation: On your marks...
Setting the grid
Let's start at the very beginning. In truth, this need not be 1 as long as the values are added in sequential order. However, this helps us check for errors by comparing to the OEIS sequences
End of explanation
"""


base = 61  # odd number please
if base % 2 == 0:
    base += 1
limit = base ** 2

"""
Explanation: Now we will build the (square) grid. All we need to determine is the size of its sides. 
Note: for the unmodified version of this problem, a $61 \times 61$ square is more than enough.
End of explanation
"""


values = [i for i in range(start, limit+1)]

"""
Explanation: All the values to be added to the grid can be generated with a list comprehension:
End of explanation
"""


grid = np.zeros((base, base), dtype=int)

"""
Explanation: Now we create a square grid with size (base, base) and all its values initialized to zero.
End of explanation
"""


midp = int(np.floor(base/2))

exes = [midp]
eyes = [midp]

grid[midp, midp] = start

"""
Explanation: Now, let's find where the midpoint of the grid is. Save these coordinates to the exes and eyes lists, and finally label the middle of the grid with the first number in the sequence.
End of explanation
"""


u = [1]
for n in range(start,limit-1):
    ulam = int(((np.floor(np.sqrt(4*n + 1)) + 3) % 4) + 1)
    u.append(ulam)

"""
Explanation: Building the spiral
The sequence A063826, indicates the steps (ESWN) one needs to make in order to build an integer spiral. By definition, $a(1) = 1$. The $n+1$-th term can be calculated as follows:
$$
a(n+1) = \left(\lfloor \sqrt{4n + 1} \rfloor + 3 \right) \bmod 4  + 1
$$
We will build this sequence with the variable ulam and appending it to a list called simply u
End of explanation
"""


for value, direction in enumerate(u):
    move_where(direction)
    start += 1
    grid[eyes[-1], exes[-1]] = start

"""
Explanation: Now, to build the spiral we just need to:

"Take a step" in the direction indicated in u,
Label the square with the values starting with start
Rinse and repeat
End of explanation
"""


# fracture points: points of no escape and their coordinates
fractures = [0]

# stop points, in order. Use these to "extend" the original 
# trapped knight problem
#fractures = [2084, 2720, 3325, 3753, 7776, 5632, 7411, 8562, 14076, 8469,
#             9231]

"""
Explanation: Get set...
Spoiler alert: the sequence is finite, hence the name of this notebook. Once the Knight gets to the square labeled 2084, there's no unvisited space left. One of the comments in the Numberphile video asked what would happen if the knight preemptively knew about the trap (more on this later) and skipped it. In order to expand on this idea, we can create a list of such traps so that the Knight never visits them. For the unmodified version of the problem, this list only contains square zero so that it stops appropriately
End of explanation
"""


visited = [1]

path_x = [midp]
path_y = [midp]

"""
Explanation: All we need to do now is to mark all visited squares and to follow the Knight on his path
End of explanation
"""


while path_x[-1] < (base - 2) and path_y[-1] < (base - 2):
    if visited[-1] != fractures[-1]:
        pos = [path_y[-1], path_x[-1]]
        next_pos = [0, 0]
        next_pos[0], next_pos[1] = get_next(pos, grid)
        if next_pos[0] > 0:
            path_x.append(next_pos[0])
            path_y.append(next_pos[1])
        else:
            break
    else:
        break

"""
Explanation: Go!
Now the main loop runs. We first check if the Knight gets too close to the edge of the grid (which shouldn't happen in the idealized infinite grid, but must be accounted for in this exercise). Then, we check that the last visited square is not one of the "trap" squares.
If we pass these checks, then the Knight's position will be updated, the next position will be calculated with get_next(pos, grid) and passed to the Knight's path_x, path_y variables.
End of explanation
"""


path_x = np.asarray(path_x)
path_y = np.asarray(path_y)

"""
Explanation: Visualization
In order to better visualize the Knight's path, we can now convert path_x, path_y to numpy arrays. The main reason for doing it now rather than before is to save memory (since np.append() creates a copy of the array every time it's called).
End of explanation
"""


%matplotlib inline
plt.figure(figsize=(8, 8))
plt.title('Figure 1: The final step')
plt.imshow(grid, alpha=0)
plt.plot(path_x, path_y, alpha=0.3, linewidth=1)
plt.scatter(path_x[-1], path_y[-1], marker='x', color='r')
plot_moves_from_number(2084)
plt.xlim(path_x[-1] - 4, path_x[-1] + 4)
plt.ylim(path_y[-1] - 4, path_y[-1] + 4)
plt.show()

"""
Explanation: Now, let's see what happened at the last step of the Knight's path
End of explanation
"""


dydx = np.linspace(0, 1, num=len(path_x))

# Switch this to -1 to cheat and "flip" the y axis so that the origin
# appears to be on the top left (image-style coordinates)
path_y1 = path_y * 1
points = np.array([path_x, path_y1]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

# Figure 2
fig = plt.figure()
axs = fig.subplots()
axs.set_aspect('equal')  # Force square aspect ratio

# Create a continuous norm to map from data points to colors
norm = plt.Normalize(dydx.min(), dydx.max())
lc = LineCollection(segments, cmap='viridis', norm=norm)
# Set the values used for colormapping
lc.set_array(dydx)
lc.set_linewidth(1)
line = axs.add_collection(lc)
axs.scatter(path_x[-1], path_y1[-1], marker='x', color='r')


# Shazam!
# =======
plt.show()

"""
Explanation: The whole path
My goal was to plot the whole path using a single colormap from matplotlib. Unfortunately, I couldn't figure this one out quickly and in depth, as I would like. Instead, I took this example out from the matplotlib gallery of examples. To quote them:

Create a set of line segments so that we can color them individually. This creates the points as a N x 1 x 2 array so that we can stack points together easily to get the segments. The segments array for line collection needs to be (numlines) x (points per line) x 2 (for x and y)
End of explanation
"""


