"""This module provides graph theory functionality."""

from heapq import heapify, heappop, heappush

def dijkstra(nodes, edges, startNode, directed):
    """Finds the length between each node in the graph and the startNode.
    Arguments:
        nodes - the set of nodes in the graph.
        edges - the set of edges in the graph. Each edge should be a 3-tuple
                containing the source and destination nodes, as well as the
                (non-negative) weight of that edge (in that order).
        startNode - the starting node for the search.
        directed - boolean, should the graph be treated as directed,
                   or instead indirected?
    Returns: a dict whose keys are nodes and whose values are the smallest cost
    to get from startNode to that particular node.
    
    E.g.
    # example: from wikipedia    
    nodes = ["1", "2", "3", "4", "5", "6"]
    edges = [("1", "2", 7), ("1", "3", 9), ("1", "6", 14), ("2", "3", 10),
             ("2", "4", 15), ("3", "4", 11), ("3", "6", 2), ("4", "5", 6),
             ("5", "6", 9)]
    d = dijkstra(set(nodes), set(edges), "1", True)
    print d
    """
    #construct a dict holding each nodes' neighbours and the cost to them
    neighbours = dict([(node, []) for node in nodes])
    for edge in edges:
        neighbours[edge[0]].append((edge[1], edge[2]))
        if not directed:
            

    
    #to every node assign a distance (starting with +inf and zero for startNode)
    distance = dict([(node, float("inf")) for node in nodes])
    distance[startNode] = 0
    
    #mark every node as unvisited
    visited = dict([(node, False) for node in nodes])
    
    #main part of algorithm
    unvisitedQ = [(0, startNode)]
    unvisitedQSet = set([startNode])
    while len(unvisitedQ) != 0:
        currentNode = heapq.heappop(unvisitedQ)[1]
        unvisitedQSet.remove(currentNode)
        for (node, edgeWeight) in neighbours[currentNode]:
            if not visited[node]:
                if distance[currentNode] + edgeWeight < distance[node]:
                    distance[node] = distance[currentNode] + edgeWeight
                if node not in unvisitedQSet:
                    heapq.heappush(unvisitedQ, (distance[node], node))
                    unvisitedQSet.add(node)
        visited[currentNode] = True
    return distance


def prim(adj_list, start_v):
    """Finds a minimal spanning tree given a graph's adjacency list. The list
    should be a dictionary whose keys are vertices. The values should be lists
    of 2-tuples. Each 2-tuple should contain a 'to vertex' and a weight.
    
    Returned is a list of edges in the minimal spanning tree, each a 3-tuple
    containing the 'to vertex', 'from vertex', and weight.
    
    E.g. from wikipedia:
    
    a = {'A': [('B', 7), ('D', 5)],
         'B': [('A', 7), ('C', 8), ('D', 9), ('E', 7)],
         'C': [('B', 8), ('E', 5)],
         'D': [('A', 5), ('B', 9), ('E', 15), ('F', 6)],
         'E': [('B', 7), ('C', 5), ('D', 15), ('F', 8), ('G', 9)],
         'F': [('D', 6), ('E', 8), ('G', 11)],
         'G': [('E', 9), ('F', 11)]}
    for from_v, to_v, weight in prim(a, 'A'):
        print from_v, to_v, weight
    """
    n = len(adj_list)
    v_new = [start_v]
    e_new = []
    q = [(weight, start_v, to_v) for to_v, weight in adj_list[start_v]]
    heapify(q)
    while len(v_new) != n:
        while True:
            weight, from_v, to_v = heappop(q)
            if from_v not in v_new or to_v not in v_new:
                break
        v_new.append(to_v)
        e_new.append((from_v, to_v, weight))
        for next_to_v, weight in adj_list[to_v]:
            if next_to_v not in v_new:        
                heappush(q, (weight, to_v, next_to_v))
    return e_new

neighbours[edge[1]].append((edge[0], edge[2]))