"""
 prims.py

 An implementation of Prim's algorithm

 Jim Mahoney | March 2013 | MIT License
"""

from weighted_graph import WeightedGraph
from data_structures import MinHeap

def prims(graph, start = None, verbose=False):
    """ Return new minimal spanning tree graph """
    if start == None:
        start = graph.vertices()[0]
    min_spanning = WeightedGraph().add_vertex(start)
    all_edges = graph.edges()
    while len(min_spanning) < len(graph):
        shortest_distance = float('inf')
        shortest_edge = ()
        for edge in all_edges:
            if min_spanning.has_one_vertex(edge):
                if edge[2] < shortest_distance:
                    shortest_distance = edge[2]
                    shortest_edge = edge
        min_spanning.add_edge(*shortest_edge)
        if verbose:
            print " end of loop mst = " + str(min_spanning.vertices())
    return min_spanning

original= WeightedGraph()
for e in (('A', 'B', 12), ('A', 'C', 5), ('A', 'D', 7),
          ('D', 'B',  4), ('D', 'C', 9), ('D', 'E', 3),
          ('D', 'F',  4), ('C', 'F', 7), ('B', 'E', 7),
          ('F', 'E',  2), ('F', 'G', 5), ('G', 'E', 2)):
    original.add_edge(*e)

print "Original tree is ", original.edges()
print "Running Prim's algorithm... "
minimal = prims(original)
print "Done. "
print "Minimal tree is ", minimal.edges()
print "Creating png files."
original.graphviz_png('minimal_span', minimal)

if __name__ == '__main__':
   import doctest
   doctest.testmod()