from graphs import MutableGraph, Directed, Undirected

class DictOfLists(MutableGraph):
    """Graphs, represented using a dict-of-lists representation.

    Only the `1' value is accepted as a valid edge label.
    For example, a graph with:
     - nodes {0, 1, 2, 3, 4, 5}
     - edges {(0, 1), (0, 2), (1, 2), (1, 3), (2, 3), (3, 2), (4, 5), (5, 2)}
    and all edge labels equal to 1 is represented as
    {0: [1, 2],
     1: [2, 3],
     2: [3],
     3: [2],
     4: [5],
     5: [2]}

    This representation can be accessed via the `data' attribute.

    For undirected dictionaries, edges are represented in both of nodes'
    adjacency lists.

    Reference: Python Patterns - Implementing Graphs by Guido Van Rossum
    http://www.python.org/doc/essays/graphs.html
    """

    def __init__(self, data=None):
        if data is None:
            data = {}
        self.nodes = self.NodesView(data)
        self.edges = self.EdgesView(data)
        self.data = data

    def neighbors(self, node):
        return {n: 1 for n in self.data[node]}

    def incoming(self, node):
        return {n: 1 for n, adj in self.data.items() if node in adj}

    class NodesView(MutableGraph.NodesView):

        def __init__(self, data):
            self.data = data

        def add(self, node):
            self.data.setdefault(node, {})

        def discard(self, node):
            data = self.data
            if node in data:
                del data[node]
                for nlist in data.values():
                    try:
                        idx = nlist.index(node)
                    except ValueError:
                        pass
                    else:
                        nlist[idx] = nlist[-1]
                        nlist.pop()

        def __len__(self):
            return len(self.data)

        def __iter__(self):
            return iter(self.data)

        def __contains__(self, node):
            return node in self.data

        @classmethod
        def _from_iterable(cls, it):
            # Needed by the Set ABC
            return {node: [] for node in it}

        # This needs to be rewritten since the stock
        # MutableMapping implementations would raise an error when
        # trying to modify the dictionary while iterating on them.
        def __iand__(self, c):
            for node in self - c:
                self.discard(node)
            return self

    class EdgesView(MutableGraph.EdgesView):

        def __init__(self, data):
            self.data = data

        def __contains__(self, edge):
            n1, n2 = edge
            return n1 in self.data and n2 in self.data[n1]

        def __getitem__(self, edge):
            n1, n2 = edge
            if n2 in self.data[n1]:
                return 1
            else:
                raise KeyError(edge)

        def __delitem__(self, edge):
            n1, n2 = edge
            nlist = self.data[n1]
            try:
                idx = nlist.index(n2)
            except ValueError:
                raise KeyError(n2)
            else:
                nlist[idx] = nlist[-1]
                nlist.pop()


class Directed_DictOfLists(Directed, DictOfLists):

    class EdgesView(DictOfLists.EdgesView):

        def __len__(self):
            return sum(map(len, self.data))

        def __setitem__(self, edge, value):
            n1, n2 = edge
            if value != 1:
                raise ValueError("Dicts of lists only have edges with value 1")
            nlist = self.data.setdefault(n1, [])
            if n2 not in nlist:
                nlist.append(n2)
                self.data.setdefault(n2, [])

        def __iter__(self):
            return ((n1, n2) for n1, nlist in self.data.items() for n2 in nlist)

class Undirected_DictOfLists(Undirected, DictOfLists):

    class NodesView(DictOfLists.NodesView):

        def discard(self, node):
            # Optimization for speed
            data = self.data
            if node in data:
                for n in data[node]:
                    data[n].remove(node)
                del data[node]

    class EdgesView(DictOfLists.EdgesView):

        def __len__(self):
            data = self.data
            selfloops = sum(node in nlist for node, nlist in data.items())
            totedges = sum(map(len, data.values()))
            return (totedges + selfloops) // 2

        def __setitem__(self, edge, value):
            n1, n2 = edge
            if value != 1:
                raise ValueError("Dicts of lists only have edges with value 1")
            data = self.data
            n1list, n2list = data.setdefault(n1, []), data.setdefault(n2, [])
            if len(n1list) <= len(n2list):
                if n2 in n1list:
                    return
            else:
                if n1 in n2list:
                    return
            n1list.append(n2)
            n2list.append(n1)

        def __delitem__(self, edge):
            super().__delitem__(edge)
            n1, n2 = edge
            super().__delitem__((n2, n1))

        def __iter__(self):
            seen = set()
            for n1, nlist in self.data.items():
                for n2 in nlist:
                    if n2 not in seen:
                        yield frozenset((n1, n2))
                seen.add(n1)

if __name__ == '__main__':
    ug = Undirected_DictOfLists()
    dg = Directed_DictOfLists()
    assert type(dg.edges) == Directed_DictOfLists.EdgesView
    ug.edges[1, 2] = 1
    ug.edges.add(2, 3)
    assert ug.edges[3, 2] == 1
    ug.edges.discard(3,2)
    assert 2, 3 not in ug
    dg.edges[1, 2] = 1
    assert dg.data == {1: [2], 2: []}
    assert (2, 1) not in dg.edges
    assert (2, 1) in ug.edges
    emptygraph = Undirected_DictOfLists()
    assert emptygraph != ug
    assert emptygraph.nodes <= ug.nodes
    ug.nodes &= {1, 5}
    assert set(ug.nodes) == {1}
    ug.nodes.clear()
    assert emptygraph.nodes <= ug.nodes
    assert emptygraph == ug