"""
 weighted_graphs.py

 A weighted graph implemented 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)):
   ...     wg.add_edge(*e)           # Note: same as add_edge(e[0], e[1], e[2])
   >>> wg.graphviz_png('weighted')   # Create weighted.png and weighted.dot .
   >>> 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']
   >>> wg.edges()[0:4]
   [('A', 'B', 12), ('A', 'C', 5), ('A', 'D', 7), ('B', 'D', 4)]
   >>> len(wg.edges())
   12
   >>> len(wg)
   7

 Jim Mahoney | Feb 2013 | MIT License
"""

class Stack(object):
    """ first-in-last-out collection
        >>> s = Stack()
        >>> s.push('a')
        >>> s.push('b')
        >>> len(s)
        2
        >>> s.pop()
        'b'
        >>> s.pop()
        'a'
        >>> len(s)
        0
    """
    def __init__(self):
        self.data = []
    def __len__(self):
        return len(self.data)
    def push(self, item):
        self.data.append(item)       # put onto right end
    def pop(self):
        return self.data.pop()       # pull from right end

class Queue(object):
    """ first-in-first-out collection
        >>> q = Queue()
        >>> q.push('a')
        >>> q.push('b')
        >>> len(q)
        2
        >>> q.pop()
        'a'
        >>> q.pop()
        'b'
        >>> len(q)
        0
    """
    def __init__(self):
        self.data = []
    def __len__(self):
        return len(self.data)
    def push(self, item):
        self.data.append(item)       # put onto right end
    def pop(self):
        return self.data.pop(0)      # pull from left end


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 add_vertex(self, vertex):
        """ Put vertex with given name into this graph. """
        if not self.has_vertex(vertex):
           self.adjacency[vertex] = []

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

    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})

    def edges(self):
        """ Return list of edges as tuples (start, end, weight) """
        result = []
        for begin in self.vertices():
            for edge_dict in self.adjacency[begin]:
                end = edge_dict['end']
                if begin < end:
                    result.append((begin, end, 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):
        """ Return text description consistent with 'dot' in graphviz """
        result = "graph {\n"
        for v in self.vertices():
            for e in filter(lambda x: x['end'] > v, self.adjacency[v]):
                result += '  %s -- %s [label="%s"]; \n' % \
                          (str(v), str(e['end']), str(e['weight']))
        return result + '  rankdir=LR;\n}\n'

    def graphviz_png(self, filename='graph'):
        """ 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())
        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()