1'''
2altgraph.GraphUtil - Utility classes and functions
3==================================================
4'''
5
6import random
7from collections import deque
8from altgraph import Graph
9from altgraph import GraphError
10
11def generate_random_graph(node_num, edge_num, self_loops=False, multi_edges=False):
12    '''
13    Generates and returns a :py:class:`~altgraph.Graph.Graph` instance with *node_num* nodes
14    randomly connected by *edge_num* edges.
15    '''
16    g = Graph.Graph()
17
18    if not multi_edges:
19        if self_loops:
20            max_edges = node_num * node_num
21        else:
22            max_edges = node_num * (node_num-1)
23
24        if edge_num > max_edges:
25            raise GraphError("inconsistent arguments to 'generate_random_graph'")
26
27    nodes = range(node_num)
28
29    for node in nodes:
30        g.add_node(node)
31
32    while 1:
33        head = random.choice(nodes)
34        tail = random.choice(nodes)
35
36        # loop defense
37        if head == tail and not self_loops:
38            continue
39
40        # multiple edge defense
41        if g.edge_by_node(head,tail) is not None and not multi_edges:
42            continue
43
44        # add the edge
45        g.add_edge(head, tail)
46        if g.number_of_edges() >= edge_num:
47            break
48
49    return g
50
51def generate_scale_free_graph(steps, growth_num, self_loops=False, multi_edges=False):
52    '''
53    Generates and returns a :py:class:`~altgraph.Graph.Graph` instance that will have *steps* \* *growth_num* nodes
54    and a scale free (powerlaw) connectivity. Starting with a fully connected graph with *growth_num* nodes
55    at every step *growth_num* nodes are added to the graph and are connected to existing nodes with
56    a probability proportional to the degree of these existing nodes.
57    '''
58    # FIXME: The code doesn't seem to do what the documentation claims.
59    graph = Graph.Graph()
60
61    # initialize the graph
62    store = []
63    for i in range(growth_num):
64        #store   += [ i ] * (growth_num - 1)
65        for j in range(i + 1, growth_num):
66            store.append(i)
67            store.append(j)
68            graph.add_edge(i,j)
69
70    # generate
71    for node in range(growth_num, steps * growth_num):
72        graph.add_node(node)
73        while ( graph.out_degree(node) < growth_num ):
74            nbr = random.choice(store)
75
76            # loop defense
77            if node == nbr and not self_loops:
78                continue
79
80            # multi edge defense
81            if graph.edge_by_node(node, nbr) and not multi_edges:
82                continue
83
84            graph.add_edge(node, nbr)
85
86
87        for nbr in graph.out_nbrs(node):
88            store.append(node)
89            store.append(nbr)
90
91    return graph
92
93def filter_stack(graph, head, filters):
94    """
95    Perform a walk in a depth-first order starting
96    at *head*.
97
98    Returns (visited, removes, orphans).
99
100    * visited: the set of visited nodes
101    * removes: the list of nodes where the node
102      data does not all *filters*
103    * orphans: tuples of (last_good, node),
104      where node is not in removes, is directly
105      reachable from a node in *removes* and
106      *last_good* is the closest upstream node that is not
107      in *removes*.
108    """
109
110    visited, removes, orphans = set([head]), set(), set()
111    stack = deque([(head, head)])
112    get_data = graph.node_data
113    get_edges = graph.out_edges
114    get_tail = graph.tail
115
116    while stack:
117        last_good, node = stack.pop()
118        data = get_data(node)
119        if data is not None:
120            for filtfunc in filters:
121                if not filtfunc(data):
122                    removes.add(node)
123                    break
124            else:
125                last_good = node
126        for edge in get_edges(node):
127            tail = get_tail(edge)
128            if last_good is not node:
129                orphans.add((last_good, tail))
130            if tail not in visited:
131                visited.add(tail)
132                stack.append((last_good, tail))
133
134    orphans = [(last_good, tail) for (last_good, tail) in orphans if tail not in removes]
135    #orphans.sort()
136
137    return visited, removes, orphans
138