"""
 prims.py

 in class

"""

from weighted_graphs import *

def one_end_in(graph, edge):
    """ Return true if one end of edge is in graph, other isn't
        >>> g = WeightedGraph()
        >>> g.add_edge('A', 'B', 1)
        >>> one_end_in(g, ('A', 'Z', 1))
        True
        >>> one_end_in(g, ('Y', 'Z', 1))
        False
        >>> one_end_in(g, ('A', 'B', 1))
        False
    """
    return (graph.has_vertex(edge[0]) and not graph.has_vertex(edge[1])) \
           or \
           (graph.has_vertex(edge[1]) and not graph.has_vertex(edge[0]))


def prims(graph):
    """ Return new minimal spanning tree graph """
    all_vertices = graph.vertices()
    all_edges = graph.edges()
    mst = WeightedGraph()
    mst.add_vertex(all_vertices[0])
    while len(mst) < len(graph):
        # print " starting loop mst = " + str(mst.vertices)
        # find smallest edge with 
        # one in in mst, other end not in mst
        shortest_distance = float('inf') # infinity
        shortest_edge = ()
        for edge in all_edges:
            # print "  looking at " + str(edge)
            # print "  shortest is " + str(shortest_distance) + " , " +  str(shortest_edge)
            if one_end_in(mst, edge):
                # print "   yes, it's in mst"
                if edge[2] < shortest_distance:
                    # print "   it's the shortest so far "
                    shortest_distance = edge[2]
                    shortest_edge = edge
        mst.add_edge(*shortest_edge)
        print " end of loop mst = " + str(mst.vertices())

g = 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)):
    g.add_edge(*e)

print "The tree is "
print g.edges()
print "Running Prim's algorithm: "
prims(g)

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