"""
 floyd.py

 An example of the Floyd-Warshall Algorithm in python.

 See https://en.wikipedia.org/wiki/Floyd%E2%80%93Warshall_algorithm .

 We number the vertices from 1 to n, and define 

    shortestPath(i, j, k) = shortest path from vertex i to vertex j 
                            using only vertices from the set {1,2,...,k}
                            between i and j.

 This means that the paths considered are for the first few values of k :
    shortestPath(i, j, 0)    i--j           # these are just the weights
    shortestPath(i, j, 1)    i--j, i--1--j
    shortestPath(i, j, 2)    i--j, i--1--j, i--2--j, i--1--2--j, i--2--1--j
    
 The problem of the shortest path is solved when k=n, in other words when
 we can include all other vertices as we travel from i to j.

 And this recursion relation lets us find each shortestPath(i, j, k)
 once we already know shortestPath(i, j, k-1) :

 Either:
          i - ... -- j          without k, using only 1,2,3,...,k-1

          i -... - k -...- j    with k, using 1,2,3,...,k-1 twice, 
                                first from i to k, then from k to j.

 We want to shorter of these two options, which is :

    shortestPath(i, j, k) = min(
       shortestPath(i, j, k-1),                           # without vertex k
       shortestPath(i, k, k-1) + shortestPath(k, j, k-1)  # with vertex k
    )

 Applying this formula repeatedly lets us calculate shortestPath(i,j)
 for each value of k, starting at k=0.

 Since this is three loops, one inside the other, each 1..n, 
 the efficiency is O(n**3) = O(|V|**3).

 Jim Mahoney | cs.bennington.college | MIT License | April 2021
"""
from random import randint
import numpy as np              # numerical python

# I want fast space-efficient numerical matrices,
# and so I'll use numpy's dtype=int32 (32bit integers)
# The largest integer of that form I'll call infinity.
# We can get those sizes from numpy :
#    np.iinfo(np.int64).max  9223372036854775807 = 2**63 - 1
#    np.iinfo(np.int32).max           2147483647 = 2**31 - 1 <== enough for me
#    np.iinfo(np.int16).max                32767 = 2**15 - 1
integer_type = np.int32
# ... but doing arithmetic on the max value gives overflow errors
#infinity = np.iinfo(integer_type).max
infinity = 1000000

def get_2D_array(n, value=0):
    """ return 2D integer array a[i][j]=a[i,j]=value,  0<=i<=n 0<=j<=n 
        >>> a = get_2D_array(n=3, value=1)
        >>> a       # a[row, col] for 0 <= row <= 3 , 0 <= col <= 3 , all 1's
        array([[1, 1, 1, 1],
               [1, 1, 1, 1],
               [1, 1, 1, 1],
               [1, 1, 1, 1]], dtype=int32)
        >>> a[1,2] = 7
        >>> a
        array([[1, 1, 1, 1],
               [1, 1, 7, 1],
               [1, 1, 1, 1],
               [1, 1, 1, 1]], dtype=int32)
    """
    return np.full((n+1, n+1), value, dtype=integer_type)

# shorthand for letters for small graph example definitions.
A='A'; B='B'; C='C'; D='D'; E='E'; F='F'; G='G'; H='H'

# A sample graph from the example at
# https://en.wikipedia.org/wiki/Dijkstra%27s_algorithm
#  1      2      3      4      5      6   wikipedia's vertex labels
edges_1 = [ (A, B, 7),
            (A, C, 9),
            (A, F, 14),
            (B, C, 10),
            (B, D, 15),
            (C, F, 2),
            (C, D, 11),
            (D, E, 6),
            (E, F, 9) ]

# another example graph :
#      A - B - E
#      |\ /| 
#      | C |            
#      |   |
#      D---+
edges_2 = [(A, B, 10),
           (A, C, 2),
           (A, D, 3),
           (B, C, 4),
           (B, D, 5),
           (B, E, 6)]

class Graph(object):
    """ A weighted undirected graph 

        Each vertex is labeled with a string, its name.
        Each edge is a triple (vertex1, vertex2, weight).

        The API is
          g = Graph(edges)              create one (edges list is optional)
          g.add_edge(v1, v2, weight)    add an edge
          g.vertices()                  return list of vertex names
          g.edges()                     return list of (v1, v2, weight)
          g.neighbors(v1)               return list of neighbors of vertex
          g.get_weight(v1, v2)          return weight of edge from v1 to v2

        # doctests : documentation & tests
        >>> g = Graph()
        >>> g.add_edge('a', 'b', 1)
        >>> g.add_edge('b', 'c', 2)
        >>> g.add_edge('c', 'd', 3)
        >>> g.add_edge('a', 'c', 4)
        >>> g.vertices()               # list of vertex names
        ['a', 'b', 'c', 'd']
        >>> g.get_weight('a', 'a')     # distance to itself is zero
        0
        >>> g.get_weight('c', 'b')     # weight of c--b is same as b--c is 2.
        2
        >>> g.get_weight('a', 'd', missing=-1)  # return -1 if no such edge
        -1
        >>> g.edges()[0:3]    # first few edges
        [('a', 'b', 1), ('a', 'c', 4), ('b', 'a', 1)]

    """
    # store vertices, edges, and weights
    # as an adjacency "dictionary of dictionaries", namely 
    # { 'vertex' : {'neighbor': weight, ...}, ...}
    def __init__(self, edges=[]):
        self._vertices = {}   # i.e. {'A':{'B':1}, 'B':{'A':1}} is A--B
        for (v1, v2, weight) in edges:
            self.add_edge(v1, v2, weight)
    def add_edge(self, vertex1, vertex2, weight=1):
        for (v1, v2) in ((vertex1, vertex2), (vertex2, vertex1)):
            if not v1 in self._vertices:
                self._vertices[v1] = {}
            self._vertices[v1][v2] = weight
    def vertices(self):
        """ return list of vertices """
        return list(self._vertices.keys())
    def neighbors(self, vertex):
        """ return list of neighbors of a given vertex """
        return list(self._vertices[vertex].keys())
    def edges(self):
        """ return edges and weights as list of (v1,v2,weight) triples"""
        result = []
        for vertex in self.vertices():
            for neighbor in self.neighbors(vertex):
                weight = self.get_weight(vertex, neighbor)
                result.append((vertex, neighbor, weight))
        return result
    def get_weight(self, vertex1, vertex2, missing=None):
        """ Return weight of edge from vertex1 to vertex2 """
        if vertex1 == vertex2: return 0    # 0 distance from vertex to itself
        try:
            return self._vertices[vertex1][vertex2]
        except:
            # return default value if that (vertex1, vertex2) isn't an edge
            return missing
    def graphviz(self, path=None):
        """ return graphviz ("dot" utility) representation """
        # The output of this can be pasted into the graphviz tool at
        # https://cs.marlboro.college/tools/graphviz/dot.cgi
        result = 'graph {\n'
        for (v1, v2, weight) in self.edges():
            thick=''
            if v1 < v2:  # only include each edge once
                if path and \
                   v1 in path and v2 in path and \
                   abs(path.index(v1) - path.index(v2)) == 1:
                    thick = ', penwidth=3'
                result += f' "{v1}" -- "{v2}" [label={weight}{thick}];\n'
        result += '}\n'
        return result
        
def floyd(graph, start, end, verbose=False):
    """ Return shortest path in graph from start to end """
    # This is an adaptation of the code at
    # https://en.wikipedia.org/wiki/Floyd%E2%80%93Warshall_algorithm .
    #
    # As in the wikipedia code, I will number the vertices from 1 to
    # n, not from 0 to n-1 as is the usual list and matrix convention
    # - that will keep the formulas and definitions the same. For both
    # the arrays and matrices, I'll use numpy data structures and
    # ignore the values at index 0, using array[i] and matrix[i,
    # j]=matrix[i][j] where 1 <= i,j <= n. (Note that numpy 2D arrays
    # can use either the python matrix[row][column] syntax, or the
    # more math-typical matrix[row, column] notation.)
    #
    #  V is the set of vertices of the graph.
    #     The Graph object labels each vertex with a string name 'A', 'B', ...
    #     In this function, we'll use numerical indeces 1, 2, ..., n
    #  n = |V| = size of the set of vertices, numbered 1 to n.
    #  dist[i, j] is a 2D array of distances from vertex i to vertex j
    #
    # This is essentially the "FloydWarshallWithPathReconstruction"
    # code from the wikipedia article.
    #
    # -- initialize the matrices 
    names = [''] + graph.vertices()  # vertex names for indices 1..n (ignore 0)
    n = len(names) - 1   # number of vertices
    distances = get_2D_array(n, value=infinity)   # dist[u, v]
    path_next = get_2D_array(n, value=0)          # next[u, v]
    for u in range(1, 1+n):
        for v in range(1, 1+n):
            distances[u, v] = graph.get_weight(names[u], names[v],
                                               missing=infinity)
            path_next[u, v] = v
    # -- apply the recursive formula for k from 1 to n.
    for k in range(1, 1+n):
        for i in range(1, 1+n):
            for j in range(1, 1+n):
                without_k = distances[i, j]                 # i ... j
                with_k = distances[i, k] + distances[k, j]  # i ... k ... j
                if with_k < without_k:                      # better with k?
                    distances[i, j] = with_k                #   yes; update.
                    path_next[i, j] = path_next[i, k]       #   remember path
        if verbose:
            print(f"--- k={k} ---")
            print("distances:")
            print(distances)
            print("path_next:")
            print(path_next)
            

    # -- reconstruct the path from start to end
    (i_here, i_end) = (names.index(start), names.index(end))
    if path_next[i_here, i_end] == 0: return []      # no path found
    path = [names[i_here]]
    while i_here != i_end:
        i_here = path_next[i_here, i_end]
        path.append(names[i_here])
    return path

# --- 2D grid - from my mazey.py -----------

def label(xy): return f"{xy[0]},{xy[1]}"
def add(xy1, xy2): return (xy1[0]+xy2[0], xy1[1]+xy2[1])  # vector addition
def scale(s, xy): return (s*xy[0], s*xy[1])               # scalar multiply
def in_grid(xy, n): return xy[0]>=0 and xy[1]>=0 and xy[0]<=n and xy[1]<=n
offsets = ((0,1), (0,-1), (1,0), (-1,0))          # four directions

def get_grid(n):
    """ return tuple of (x,y) grid points, (0,0) <= (x,y) <= (n,n) """
    return tuple((x,y) for x in range(n+1) for y in range(n+1))

def get_neighbor_points(xy, n):
    """ return points adjacent to xy within the 0 to (n-1) grid """
    # For example, neighbors of (0,0) are [(0,1), (1,0)]
    # At the edges of the grid, points have three neighbors.
    # At the corners, they have two neighbors.
    # And elsewhere each has four neighbors.
    neighbors = []
    for offset in offsets:
        neighbor = add(offset, xy)
        if in_grid(neighbor, n):
            neighbors.append(neighbor)
    return neighbors

def in_wall(xy, n):
    # . . . . . . . . . .    # n = 10 x 10 grid
    # . . . . . . . . . .
    # . . . . . . . . . .
    # . . X X X X X X . .
    # . . X X X X X X . .    # the wall are the X's :
    # . . X X X X X X . .    # n//2 +-1 in y, n//4 to 3*n//4 in x
    # . . . . . . . . . .    
    # . . . . . . . . . .    
    # . . . . . . . . . .
    (ymin, ymax) = (n//2 - 1, n//2 + 1)
    (xmin, xmax) = (n//4, 3*n//4)
    return xmin <= xy[0] <= xmax and ymin <= xy[1] <= ymax

def make_graph_grid(n=10):
    graph = Graph()
    vertices = get_grid(n)
    for xy1 in vertices:
        for xy2 in get_neighbor_points(xy1, n):
            distance = randint(n//2, n//2 + n//4)
            if not (in_wall(xy1, n) or in_wall(xy2, n)):
                graph.add_edge(label(xy1), label(xy2), distance)
    return graph

def analyze(graph, name, start, end, verbose=False):
    """ run Dijkstra's algorithm on graph, print analysis """
    pass
    print('- '*30)
    print(f"Floyd's shortest path from {start} to {end} in {name} :")
    n = len(graph.vertices())
    print(f"graph has {n} vertices; Floyd matrix has {n*n} entries.")
    print()
    #print(" vertices: ", graph.vertices())
    #print(" start is in vertices: ", start in graph.vertices())
    #print(" end is in vertices: ", end in graph.vertices())
    #print()
    path = floyd(graph, start, end, verbose=verbose)
    print(path)
    print()
    print(graph.graphviz(path))
    print()

def main():

    analyze(Graph(edges_1), 'example 1', A, E, verbose=True)
    analyze(Graph(edges_2), 'example 2', A, E, verbose=True)

    for n in (10, 15):
        graph = make_graph_grid(n)
        (start, end) = (label((n//2, 0)), label((n//2, n)))
        analyze(graph, f'grid {n} x {n}', start, end)
    
if __name__ == '__main__':
   import doctest
   doctest.testmod()
   main()