import heapq as heap


# function which takes graph, vertex set, start and goal, uses djikstra algorithm to establish shortest paths from start
# every other vertex in graph and returns shortest path between start and goal
def dijkstra_algorithm(graph, start, goal, is_directed=False):
    # dictionary which will store for every vertex it's distance to start
    dist = dict()

    # dictionary which keeps if vertex was already popped from queue
    visited = dict()

    # dictionary which keeps for every vertex which other vertex is previous on the shortest path from start
    prev = dict()

    # list which is used to keep priority queue with vertexes to be explored by the algorithm
    queue = []

    # dictionary for keeping neighbors of vertexes.
    point_set = dict()
    for arc in graph.keys():
        point_set[arc[0]] = []
        point_set[arc[1]] = []

    # creating dictionary which for every vertex keeps all it's neighbors
    for arc in graph.keys():
        point_set[arc[0]].append(arc[1])
        if not is_directed:
            point_set[arc[1]].append(arc[0])

    # initialization of needed lists
    for vertex in point_set:
        # setting distance of every vertex from the starting vertex to infinity
        dist[vertex] = float('inf')
        # setting previous vertex in currently found shortest path for every vertex to None
        prev[vertex] = None
        # setting flag for every vertex keeping track if it was already visited my the algorithm
        visited[vertex] = False

    # pushing start vertex to priority queue
    heap.heappush(queue, (0, start))
    # setting distance to start for start vertex to 0
    dist[start] = 0
    # setting start vertex not to be added to priority queue again
    visited[start] = True

    # main loop
    while len(queue) > 0:
        # getting first vertex from the priority queue and saving it in current variable
        current = heap.heappop(queue)
        # iterating trough all neighbors of the current vertex
        for neighbor in point_set[current[1]]:
            # initializing potential distance to be replaced with the current neighbors
            new_dist = 0
            #
            if not is_directed:
                if neighbor > current[1]:
                    new_dist = dist[current[1]] + graph[(neighbor, current[1])]
                else:
                    new_dist = dist[current[1]] + graph[(current[1], neighbor)]
            else:
                new_dist = dist[current[1]] + graph[(current[1], neighbor)]
            if new_dist < dist[neighbor]:
                if dist[neighbor] == float('inf'):
                    dist[neighbor] = new_dist
                    prev[neighbor] = current[1]
                    heap.heappush(queue, (new_dist, neighbor))
                else:
                    dist[neighbor] = new_dist
                    prev[neighbor] = current[1]
                    for k in range(len(queue)):
                        if queue[k][1] == neighbor:
                            queue[k] = (new_dist, neighbor)
                            heap.heapify(queue)

    temp = goal
    shortest_path = [goal]
    while prev[temp] is not None:
        shortest_path.append(prev[temp])
        temp = prev[temp]
    shortest_path.reverse()
    return shortest_path, dist[goal]


if __name__ == "__main__":
    g = {
        (2, 1): 3,
        (3, 2): 2,
        (5, 3): 1,
        (9, 5): 5,
        (10, 9): 4,
        (9, 4): 3,
        (4, 1): 4,
        (7, 1): 6,
        (3, 1): 4,
        (6, 2): 3,
        (8, 6): 8,
        (8, 3): 2,
        (9, 1): 8
    }

    g2 = {
        (4, 1): 6,
        (1, 3): 4,
        (1, 2): 2,
        (2, 4): 5,
        (3, 4): 1,
        (5, 3): 1
    }

    print(dijkstra_algorithm(g, 7, 10))
    print(dijkstra_algorithm(g2, 1, 4, is_directed=True))
    print(dijkstra_algorithm(g2, 1, 5, is_directed=True))