"""
 weighted_graph.py

 A weighted graph implementation 
 using adjacency lists.
 
   # This sample graph is from pg 196 of Skiena's "Algorithm Design Manual".
   >>> wg = 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)):
   ...     tmp = wg.add_edge(*e)

   >>> wg.graphviz_png('weighted')   # Create weighted.png and weighted.dot
   >>> subg = WeightedGraph().add_edge('A', 'B', 12).add_edge('B', 'D',  4)
   >>> wg.graphviz_png('subgraph', subg)

   >>> wg.search(direction = 'breadth')
   ['A', 'B', 'C', 'D', 'E', 'F', 'G']

   >>> wg.search(direction = 'depth', start = 'B')
   ['B', 'E', 'G', 'F', 'C', 'D', 'A']

   >>> wg.vertices()
   ['A', 'B', 'C', 'D', 'E', 'F', 'G']
   >>> len(wg)           # Number of vertices in graph
   7
   >>> len(wg.edges())   # Number of edges in graph
   12

   >>> wg.edges()[0:4]   # First 4 edges (alphabetic order)
   [('A', 'B', 12), ('A', 'C', 5), ('A', 'D', 7), ('B', 'D', 4)]
   >>> wg.edges('B')     # Edges containing B (vertices in alphabetic order)
   [('A', 'B', 12), ('B', 'D', 4), ('B', 'E', 7)]

   >>> wg.has_vertex('A')                   # A is in the graph
   True
   >>> wg.has_one_vertex(('A', 'Z', 1))     # In the edge, A is in & Z isn't.
   True

 Jim Mahoney | March 7 2013 | MIT License
"""

from data_structures import Stack, Queue    

class WeightedGraph(object):

    def __init__(self):
        # self.adjacency is an adjacency list in this format :
        #  { vertex_begin : [ {'end': vertex_end, 'weight': xxx}, ...], ... }
        # with edges from vertex_begin to vertex_end with given weight
        self.adjacency = {} 

    def __repr__(self):
        return "<WeightedGraph with %s vertices at 0x%x>" % (len(self), id(self))

    def add_vertex(self, vertex):
        """ Put vertex with given name into this graph. """
        if not self.has_vertex(vertex):
           self.adjacency[vertex] = []
        return self

    def has_vertex(self, vertex):
        """ Return True if vertex is in the graph """
        return vertex in self.adjacency.keys()

    def has_one_vertex(self, edge):
        """ Return True if the graph contains exactly one of the 
            two vertices in the edge. """
        # This is *not* a check on whether or not the edge is in the graph.
        return (self.has_vertex(edge[0]) and not self.has_vertex(edge[1])) \
            or (self.has_vertex(edge[1]) and not self.has_vertex(edge[0]))

    def vertices(self):
        """ Return a list of the vertex names. """
        return sorted(self.adjacency.keys())

    def __len__(self):
        """ Return size of matrix via python's len() function """
        return len(self.adjacency)

    def add_edge(self, begin, end, weight=1):
        self.add_vertex(begin)
        self.add_vertex(end)
        self.adjacency[begin].append({'end': end, 'weight': weight})
        self.adjacency[end].append({'end': begin, 'weight': weight})
        return self

    def edges(self, vertex=None):
        """ Return list of edges as tuples (vertex_A, vertex_B, weight).
            (Optionally only those containing a given vertex.) """
        # Only one edge for each pair of connected vertices is
        # returned as (A, B, Weight) and always in alphabetic order
        # with A < B. Therefore, if only edges containing a given
        # vertex are returned, that vertex may be A or B.
        # The list is sorted before it is returned.
        result = []
        if vertex == None:
            vertices = self.vertices()
        else:
            vertices = (vertex)
        for begin in vertices:
            for edge_dict in self.adjacency[begin]:
                end = edge_dict['end']
                if begin < end:
                    result.append((begin, end, edge_dict['weight']))
                elif vertex != None:
                    # If only edges from a given vertex are desired,
                    # then each one should be included even when
                    # the two vertices aren't already in alphabetic order
                    result.append((end, begin, edge_dict['weight']))                    
        result.sort()
        return result

    def neighbors(self, vertex):
        """ Return list of neighbors of vertex """
        # self.adjacency[vertex] is a list of dict of {name:, distance:, ...}
        return map(lambda e: e['end'], self.adjacency[vertex])

    def _as_graphviz(self, subgraph=None):
        """ Return text description of subgraph within graph """
        if subgraph != None:
            sub_edges = set(subgraph.edges())
        else:
            sub_edges = set()
        # print " sub_edges = '%s' " % sub_edges
        result = "graph {\n"
        for v in self.vertices():
            for e in filter(lambda x: x['end'] > v, self.adjacency[v]):
                ve = (v, e['end'], e['weight'])
                # print "  ve = '%s' " % str(ve)
                if ve in sub_edges:
                    color = ', color="red"'
                else:
                    color = ''
                result += '  %s -- %s [label="%s" %s]; \n' % \
                          (str(v), str(e['end']), str(e['weight']), color)
        return result + '  rankdir=LR;\n}\n'

    def graphviz_png(self, filename='graph', subgraph=None):
        """ Create e.g. graph.dot and graph.png file with image of graph """
        from subprocess import call
        dot_filename = filename + '.dot'
        png_filename = filename + '.png'
        dot_file = open(dot_filename, 'w')
        dot_file.write(self._as_graphviz(subgraph))
        dot_file.close()
        call(['dot', '-Tpng', '-o' + png_filename, dot_filename])

    def search(self, start=None, direction='breadth', function=lambda x: x):
        """ Traverse the graph, applying the given function to each vertex.
            Return a list of the function return values, leaving out None."""
        # The traditional "brute force search"
        if len(self) == 0:
            return []
        if start == None:
            start = self.vertices()[0]
        if direction == 'breadth':
            fringe = Queue()        # breadth first search
        else:
            fringe = Stack()        # depth first search
        fringe.push(start)
        visited = {}                # Vertices that have been visited.
        results = []                # Output from functions run on vertices
        while len(fringe) > 0:
            vertex = fringe.pop()
            if vertex not in visited:
                visited[vertex] = True
                result = function(vertex)
                if result != None:
                    results.append(result)
                for child in self.neighbors(vertex):
                    fringe.push(child)
        return results



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