"""
 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 reverse(tuple):
    # GD Guido ...
    # instead of tuple.reversed()
    temp = list(tuple)
    temp.reverse()
    return temp

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)

    fringe = MinHeap()
    for edge in graph.edges(start):
        # put the edges that include the start vertex into fringe
        # order by distance, so reverse to put that first
        fringe.push(reverse(edge))

    print "fringe = " + str(fringe.data)

    while len(min_spanning) < len(graph):

        looping = True
        while looping:
            smallest_edge = reverse(fringe.pop())
            if min_spanning.has_one_vertex(smallest_edge):
                looping = False
                min_spanning.add_edge(smallest_edge)

        # Figure out which end of smallest_edge was just added.
        # Add all the neighboring edges of that to the fringe.
        
        
    #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, verbose=True)
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()