# Source: content/notes/math/discrete-math.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

from collections import deque
from graphlib import TopologicalSorter, CycleError

graph = {"a": ["b", "c"], "b": ["d"], "c": ["d"], "d": []}
def bfs(graph, start):
    distance, queue = {start: 0}, deque([start])
    while queue:
        node = queue.popleft()
        for neighbor in graph[node]:
            if neighbor not in distance:
                distance[neighbor] = distance[node]+1
                queue.append(neighbor)
    return distance

assert bfs(graph, "a") == {"a": 0, "b": 1, "c": 1, "d": 2}
assert bfs(graph, "d") == {"d": 0}
predecessors = {node: set() for node in graph}
for node, neighbors in graph.items():
    for neighbor in neighbors:
        predecessors[neighbor].add(node)
order = tuple(TopologicalSorter(predecessors).static_order())
position = {node: i for i, node in enumerate(order)}
assert all(position[u] < position[v] for u in graph for v in graph[u])
predecessors["a"].add("d")
try:
    tuple(TopologicalSorter(predecessors).static_order())
except CycleError:
    pass
else:
    raise AssertionError("Expected a directed cycle")
print("BFS:", bfs(graph, "a"), "topological order:", order)