"""
 graphs.py

 implementing some code for dealing with graphs
 see http://cs.marlboro.edu/courses/spring2013/algorithms/notes/Feb_21 and Feb_26

 Initial API idea :

    g = Graph(optional_properties)
    g.add_edge(("vertex_name_1", "vertex_name_2"))
    g.vertices()   # return list of vertices
    g.edges()      # return list of edges (each is a pair of vertices)
    g.search("starting node", breadth_or_depth, function_to_apply)

 An example graph 

      A - B - E
      |\ /| 
      | C |            
      |   |
      D---+

 Tests for Graph1, adjacency matrix version

    >>> g = Graph1()
    >>> for edge in (('A', 'B'), ('A', 'C'), ('A', 'D'),
    ...              ('B', 'C'), ('B', 'D'), ('B', 'E')):
    ...   g.add_edge(edge)
    >>> g.vertices()
    ['A', 'B', 'C', 'D', 'E']
    >>> len(g)  # number of vertices
    5
    >>> g.edges()
    [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('B', 'E')]
    >>> g.has_vertex('H')
    False
    >>> g.has_edge(('B', 'A'))   # A-B same edge as B-A
    True
    >>> g.neighbors('A')
    ['B', 'C', 'D']
    >>> g.search(direction = 'breadth')
    ['A', 'B', 'C', 'D', 'E']
    >>> g.search(direction = 'depth')
    ['A', 'D', 'B', 'E', 'C']

 Tests for Graph2, adjacency list version (implemented with python dict)

    >>> g = Graph2()
    >>> for edge in (('A', 'B'), ('A', 'C'), ('A', 'D'),
    ...              ('B', 'C'), ('B', 'D'), ('B', 'E')):
    ...   g.add_edge(edge)
    >>> g.vertices()
    ['A', 'B', 'C', 'D', 'E']
    >>> len(g)  # number of vertices
    5
    >>> g.edges()
    [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('B', 'E')]
    >>> g.has_vertex('H')
    False
    >>> g.has_edge(('B', 'A'))   # A-B same edge as B-A
    True
    >>> g.neighbors('A')
    ['B', 'C', 'D']
    >>> g.search(direction = 'breadth')
    ['A', 'B', 'C', 'D', 'E']
    >>> g.search(direction = 'depth')
    ['A', 'D', 'B', 'E', 'C']

 Graphviz
    This command 
      g.graphviz_png()
    will create graph.dot and a graph.png image of the graph.


 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 GraphAPI(object):
    """ base class for graph objects """
    # Don't instantiate this class - inherit from it to build other classes.
    # Several methods below assume that the vertices are in self.vertex_data.

    def has_vertex(self, vertex):
        """ Return true if that vertex is in the graph """
        return vertex in self.vertex_data

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

    def _as_graphviz(self):
        """ Return text description consistent with 'dot' in graphviz """
        result = "graph {\n"
        for v in self.vertices():
            smaller_neighbors = filter(lambda x: x > v, self.neighbors(v))
            if len(smaller_neighbors) > 0:
                result += '  %s -- { %s }; \n' % \
                          (str(v), ' '.join(smaller_neighbors))
        return result + '}\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."""
        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

    # --- override these in an inherited class --
    def __init__(self):
        pass
    def add_vertex(self, vertex):
        pass
    def vertices(self):
        return []
    def edges(self):
        return []
    def neighbors(self, vertex):
        return []
    def has_edge(self, vertex):
        return False

class Graph1(GraphAPI):
    """ An undirected simple graph """
    # adjacency matrix implementation with fixed maximum size.
    #
    # (Converting back and forth between vertex data and
    # the corresponding index into the adjacency matrix feels awkward.)
    #      
    def __init__(self, max_size = 100):
        self.max_size = max_size
        self.adjacency_matrix = [[0]*self.max_size
                                 for i in range(self.max_size)]
        self.vertex_data = []

    def vertices(self):
        return sorted(self.vertex_data)

    def add_vertex(self, vertex):
        """ Add vertex with given data if not over max_size. """
        if not self.has_vertex(vertex) and len(self) < self.max_size:
            self.vertex_data.append(vertex)

    def _index(self, vertex):
        """ Return index of given vertex data """
        if self.has_vertex(vertex):
            return self.vertex_data.index(vertex)
        else:
            return None
    def edges(self):
        """ Return sorted list of edges as vertex data pairs. """
        result = []
        for i in xrange(self.max_size):
            for j in xrange(i):
                if self.adjacency_matrix[i][j] == 1:
                    result.append(tuple(sorted((self.vertex_data[i],
                                                self.vertex_data[j]))))
        return sorted(result)
    def add_edge(self, edge, weight=1):
        for vertex in edge:
            self.add_vertex(vertex)
        i = self._index(edge[0])
        j = self._index(edge[1])
        if i != None and j != None:
            self.adjacency_matrix[i][j] = weight
            self.adjacency_matrix[j][i] = weight

    def has_edge(self, edge):
        i = self._index(edge[0])
        j = self._index(edge[1])
        if i != None and j != None:
            return self.adjacency_matrix[i][j] != 0
        else:
            return False

    def neighbors(self, vertex):
        """ Return a list of neighboring vertices """
        i = self._index(vertex)
        return sorted(map(lambda k: self.vertex_data[k],
                          filter(lambda j: self.adjacency_matrix[i][j] != 0,
                                 range(self.max_size))))

class Graph2(GraphAPI):
    """ An undirected simple graph """
    # adjancy list using python dict : {'vertix':[neighbors], ...}

    def __init__(self):
        self.vertex_data = {}

    def add_vertex(self, vertex):
       if not self.has_vertex(vertex):
           self.vertex_data[vertex] = []

    def vertices(self):
        """ Return sorted list of vertex data. """
        return sorted(self.vertex_data.keys())

    def edges(self):
        """ Return sorted list of edges as vertex data pairs. """
        result = []
        for v in self.vertex_data.keys():
            for n in self.neighbors(v):
                if v < n:                    # only return (a,b) in order
                    result.append((v, n))
        return sorted(result)
    
    def add_edge(self, edge, weight=1):
        for vertex in edge:
            self.add_vertex(vertex)
        if edge[1] not in self.vertex_data[edge[0]]:      # a is b's neighbor
            self.vertex_data[edge[0]].append(edge[1])
        if edge[0] not in self.vertex_data[edge[1]]:      # b is a's neighbor
            self.vertex_data[edge[1]].append(edge[0])

    def has_edge(self, edge):
        return self.vertex_data.has_key(edge[0]) and \
               edge[1] in self.vertex_data[edge[0]]

    def neighbors(self, vertex):
        """ Return a list of neighboring vertices """
        return sorted(self.vertex_data[vertex])

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