graph.py

import networkx as nx
import matplotlib.pyplot as plt
from collections import deque

def dag_levels(adj_list):
    """Return level of each node (distance from root layer)."""

    # Build indegree to find roots
    indegree = {node: 0 for node in adj_list}
    for u in adj_list:
        for v in adj_list[u]:
            indegree[v] += 1

    queue = deque([u for u in adj_list if indegree[u] == 0])
    level = {u: 0 for u in adj_list}

    while queue:
        u = queue.popleft()
        for v in adj_list[u]:
            level[v] = max(level[v], level[u] + 1)
            queue.append(v)

    return level


def visualize_smooth_dag(adj_list):
    G = nx.DiGraph()

    for u, nbrs in adj_list.items():
        for v in nbrs:
            G.add_edge(u, v)

    # Compute hierarchical levels
    levels = dag_levels(adj_list)

    # Start with a spring layout
    pos = nx.spring_layout(G, seed=42)

    # Now push nodes down by level (soft constraint)
    for node in pos:
        x, y = pos[node]
        pos[node] = (x, -levels[node])  # lock y approx by level

    # Give one more spring iteration but freeze Y positions
    for node in pos:
        pos[node] = (pos[node][0], pos[node][1])  # already set

    nx.draw(
        G,
        pos,
        with_labels=True,
        node_size=1500,
        node_color="skyblue",
        arrows=True,
        font_size=12
    )

    plt.title("Smooth Parent → Child DAG Layout")
    plt.show()


# Example DAG
adj = {
"0" : ["5","3","8"],
"1" : [],
"2" : ["0","7"],
"3" : ["4","8"],
"4" : [],
"5" : [],
"6" : [],
"7" : ["9","6"],
"8" : ["1"],
"9" : ["4"]
}

visualize_smooth_dag(adj)