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 = []

    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])

    for vertex in point_set:
        dist[vertex] = float('inf')
        prev[vertex] = None
        visited[vertex] = False

    heap.heappush(queue, (0, start))
    dist[start] = 0
    visited[start] = True

    while len(queue) > 0:
        current = heap.heappop(queue)
        for neighbor in point_set[current[1]]:
            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]:
                dist[neighbor] = new_dist
                for k in range(len(queue)):
                    if queue[k][1] == neighbor:
                        queue[k] = (new_dist, neighbor)
                        heap.heapify(queue)
                prev[neighbor] = current[1]

            if not visited[neighbor]:
                visited[neighbor] = True
                heap.heappush(queue, (new_dist, neighbor))
    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('')
    print(dijkstra_algorithm(g, 7, 10))
    print(dijkstra_algorithm(g2, 1, 5, is_directed=True))