"""
 dijkstra.py

 An example of Dijkstra's Algorithm in python.

 Jim Mahoney | cs.bennington.college | MIT License | April 2021
"""
from heapq import heappop, heappush     # python's built-in priority queue
from random import randint

infinity = float('inf')        # largest floating point number
A='A'; B='B'; C='C'; D='D'     # shortcuts for typing letters
E='E'; F='F'; G='G'; H='H'     #   for graph example defitions

# 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 HeapEtc:
    """ A "get smallest" collection which stores (name, weight) entries
        and also implements (a) a fast lookup, (b) "in" membership operation,
        and (c) decreasing the value of a name  already in the collection.

        The API is
          h = HeapEtc()
          h.push(name, value) add (or decrement) an entry, O(log n)
          h.pop()             return smallest value (name, value), O(log n)
          h.get(name)         return its value (or None if missing), O(1)
          name in h           return True if name is in the collection
          len(h)              number of entries in the collection, O(1)

        >>> h = HeapEtc()
        >>> h.push('A', 10)
        >>> h.push('B', 5)
        >>> h.push('C', 7)
        >>> len(h)
        3
        >>> h.get('B')      # B's value ?
        5
        >>> 'B' in h        # is 'B' in the heap?
        True
        >>> h.pop()
        ('B', 5)
        >>> h.push('A', 2)   # decrease A's value.
        >>> h.pop()          # smallest is A
        ('A', 2)
        >>> h.pop()
        ('C', 7)
        >>> len(h)
        0
    """

    # The implementation here uses both a python dict
    # and a python heap, along with a trick: to decrement
    # a value, it is just pushed onto the heap, leaving
    # two (name, value) entries in the heap with the same name.
    # The first time that name is popped, the smaller value
    # will come out, and it will also be removed from the dictionary.
    # If a later pop() gives the invalid original (name, value),
    # since that name isn't in the dictionary  it is simply ignored
    # and we pop() again. The data is stored in the heap as (value,name)
    # so that when sorted, smaller values come first.

    def __init__(self):
        self.heap = []  # [(value,name), (old_value,name), ...]
        self.items = {} # {name:value, ...} 
    def __len__(self):
        return len(self.items)
    def __contains__(self, name):
        # This implements a membership test "x in heap" syntax,
        # using the same syntax that python does to see
        # if name is a key in the items dict.
        return name in self.items
    def push(self, name, value):
        """ add (name,value) to the collection, or decrease name's value """
        self.items[name] = value
        heappush(self.heap, (value, name))
    def get(self, name):
        """ return name's value """
        return self.items[name]
    def pop(self):
        """ return and remove (name, value) item with the smallest value """
        # heappop() may find old no-longer-valid (name, value) pairs,
        # since updating (decreasing) a value pushes a new (name,value)
        # into the heap. Valid names are the ones still in the items dict.
        while len(self.heap) > 0:      # keep going until we find a valid one
            (value, name) = heappop(self.heap)       # get smallest
            if name in self.items:                   # valid ?
                del self.items[name]                 # yes: remove from dict
                return (name, value)                 #      and return it
        return None                    # Nothing left.

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

        The API is
          g = Graph(edges)                create one (edges is optional)
          g.add_edge('v1', 'v2', weight)  add an edge
          g.vertices()                    return list of vertices
          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

        >>> 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()
        ['a', 'b', 'c', 'd']

    """
    # 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):
        """ Return weight of edge from vertex1 to vertex2 """
        try:
            return self._vertices[vertex1][vertex2]
        except:
            # something wasn't found - illegal vertex or no such edge.
            return None
    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 dijkstra(graph, start, end):
    """ Return shortest path in graph from start to end """
    # -- initialize --
    previous = {}     # linked list keeping track of path back links
    heap = HeapEtc()  # min heap, best distance so far to each vertex
    for vertex in graph.vertices():
        previous[vertex] = None
        heap.push(vertex, infinity)
    heap.push(start, 0)
    # -- loop over each next-closest vertex
    while len(heap) > 0:
        (vertex, distance) = heap.pop()   # move to closest vertex
        if vertex == end:
            break     # found goal, so stop here and construct path
        for neighbor in graph.neighbors(vertex):
            # update distances for neighbors of this vertex,
            # if it hasn't been visited yet.
            if neighbor in heap:
                old_distance = heap.get(neighbor)
                new_distance = distance + graph.get_weight(vertex, neighbor)
                if new_distance < old_distance:
                    # found better path to this neighbor , through vertex
                    previous[neighbor] = vertex        # remember path
                    heap.push(neighbor, new_distance)  # remember distance
    # -- have answer; now find the path by tracking back from end 
    path = []                  # the answer that we want
    vertex = end               # the last vertex in the path
    while vertex != None:      # follow the linked list of previous vertices
        path.append(vertex)
        try:
            vertex = previous[vertex]
        except:
            print(previous)
            print("-- OOPS --")
            vertex = None
    path.reverse()             # flip from end->start to start->end
    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):
    """ run Dijkstra's algorithm on graph, print analysis """
    print('- '*30)
    print(f"Dijkstra's shortest path from {start} to {end} in {name} :")
    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 = dijkstra(graph, start, end)
    print(path)
    print()
    print(graph.graphviz(path))
    print()
    

def main():

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

    for n in (10, 15, 40):
        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()