Visualization

def mpl_draw(graph, pos = None, with_labels: bool = False, labels = None, node_list = None, node_size: int = 300, node_color = 'red', node_shape: str = 'o', alpha: float = 1.0, cmap = None, vmin = None, vmax = None, ax = None, linewidths = None, edgecolors = None, label = None, style: str = 'solid', width = None, edge_color = 'black', arrows: bool = True, arrowstyle: str = '-|>', arrowsize: int = 10, edge_cmap = None, edge_vmin = None, edge_vmax = None, connectionstyle: str = 'arc3', min_source_margin: int = 0, min_target_margin: int = 0, font_size: int = 12, font_color = 'black', font_weight: str = 'normal', font_family: str = 'sans-serif', bbox = None, horizontalalignment: str = 'center', verticalalignment: str = 'center', transform = None, clip_on: bool = True, node_attrs = None, edge_attrs = None):
    """
    Draw graph using matplotlib with extensive customization options.
    
    Provides comprehensive control over all visual elements including
    nodes, edges, labels, colors, and styling.
    
    Parameters:
    - graph: Input graph (PyGraph or PyDiGraph)
    - pos (dict, optional): Node positions {node_id: (x, y)}, auto-generated if None
    - with_labels (bool): Display node labels
    - labels (dict, optional): Custom node labels {node_id: label}
    - node_list (list, optional): Specific nodes to draw
    - node_size (int or list): Node sizes in points
    - node_color (color or list): Node colors
    - node_shape (str): Node shape ('o', 's', '^', etc.)
    - alpha (float): Node transparency (0-1)
    - cmap (colormap): Colormap for node colors
    - vmin, vmax (float): Color scale limits
    - ax (matplotlib.axes): Axes to draw on
    - linewidths (float or list): Node border widths
    - edgecolors (color or list): Node border colors
    - label (str): Graph label for legend
    - style (str): Edge line style ('solid', 'dashed', 'dotted')
    - width (float or list): Edge line widths
    - edge_color (color or list): Edge colors
    - arrows (bool): Draw arrows for directed graphs
    - arrowstyle (str): Arrow style ('-|>', '->', '<->', etc.)
    - arrowsize (int): Arrow size
    - edge_cmap (colormap): Colormap for edge colors
    - edge_vmin, edge_vmax (float): Edge color scale limits
    - connectionstyle (str): Edge curve style ('arc3')
    - min_source_margin, min_target_margin (int): Arrow margins
    - font_size (int): Label font size
    - font_color (color): Label text color
    - font_weight (str): Label font weight ('normal', 'bold')
    - font_family (str): Label font family
    - bbox (dict): Label bounding box properties
    - horizontalalignment (str): Label horizontal alignment
    - verticalalignment (str): Label vertical alignment
    - transform (matplotlib.transforms): Coordinate transformation
    - clip_on (bool): Clip drawing to axes bounds
    - node_attrs (callable): Function to extract node attributes
    - edge_attrs (callable): Function to extract edge attributes
    
    Returns:
    None (displays graph using matplotlib)
    """

def graphviz_draw(graph, node_attr_fn = None, edge_attr_fn = None, graph_attr = None, filename = None, image_type = None, method: str = 'dot'):
    """
    Draw graph using Graphviz with DOT language attributes.
    
    Leverages Graphviz's powerful layout engines and styling
    capabilities for professional-quality graph visualization.
    
    Parameters:
    - graph: Input graph (PyGraph or PyDiGraph)
    - node_attr_fn (callable, optional): Function returning node attributes dict
    - edge_attr_fn (callable, optional): Function returning edge attributes dict
    - graph_attr (dict, optional): Global graph attributes
    - filename (str, optional): Output file path
    - image_type (str, optional): Output format ('png', 'svg', 'pdf', etc.)
    - method (str): Graphviz layout engine ('dot', 'neato', 'fdp', 'sfdp', 'circo', 'twopi')
    
    Returns:
    Image or None: PIL Image object if no filename specified, None if saved to file
    """

import rustworkx as rx
import matplotlib.pyplot as plt

# Create sample graph
graph = rx.generators.karate_club_graph()

# Basic visualization with default settings
plt.figure(figsize=(10, 8))
rx.visualization.mpl_draw(graph, with_labels=True)
plt.title("Karate Club Graph")
plt.axis('off')
plt.show()

# Create weighted graph
weighted_graph = rx.PyGraph()
nodes = weighted_graph.add_nodes_from(['Hub', 'A', 'B', 'C', 'D'])
edges = [
    (0, 1, 0.8),  # Hub to A: strong
    (0, 2, 0.3),  # Hub to B: weak  
    (0, 3, 0.9),  # Hub to C: very strong
    (0, 4, 0.1),  # Hub to D: very weak
    (1, 2, 0.5),  # A to B: medium
    (2, 3, 0.7),  # B to C: strong
]
weighted_graph.add_edges_from(edges)

# Custom layout
pos = rx.spring_layout(weighted_graph, seed=42)

# Customize visualization based on data
node_sizes = [1000 if i == 0 else 300 for i in range(5)]  # Hub is larger
node_colors = ['red' if i == 0 else 'lightblue' for i in range(5)]

# Edge widths based on weights
edge_weights = [weighted_graph.get_edge_data(s, t) for s, t in weighted_graph.edge_list()]
edge_widths = [w * 5 for w in edge_weights]  # Scale for visibility

plt.figure(figsize=(10, 8))
rx.visualization.mpl_draw(
    weighted_graph,
    pos=pos,
    node_size=node_sizes,
    node_color=node_colors,
    width=edge_widths,
    with_labels=True,
    labels={0: 'HUB', 1: 'A', 2: 'B', 3: 'C', 4: 'D'},
    font_size=16,
    font_weight='bold'
)
plt.title("Weighted Network with Custom Styling")
plt.axis('off')
plt.show()

# Create directed graph
digraph = rx.PyDiGraph()
process_nodes = digraph.add_nodes_from(['Start', 'Process1', 'Process2', 'Decision', 'End'])
process_edges = [
    (0, 1, 'begin'),
    (1, 2, 'data'),
    (2, 3, 'evaluate'),
    (3, 4, 'success'),
    (3, 1, 'retry')  # Feedback loop
]
digraph.add_edges_from(process_edges)

# Hierarchical layout
pos = rx.shell_layout(
    digraph,
    nlist=[[0], [1, 2], [3], [4]],  # Arrange in levels
    scale=2.0
)

# Directed graph styling
plt.figure(figsize=(12, 8))
rx.visualization.mpl_draw(
    digraph,
    pos=pos,
    with_labels=True,
    labels={i: label for i, label in enumerate(['Start', 'Process1', 'Process2', 'Decision', 'End'])},
    node_color=['lightgreen', 'lightblue', 'lightblue', 'orange', 'lightcoral'],
    node_size=1500,
    arrows=True,
    arrowsize=20,
    arrowstyle='->',
    edge_color='gray',
    font_size=10,
    font_weight='bold'
)
plt.title("Process Flow Diagram")
plt.axis('off')
plt.show()

# Graph with numeric node data for colormap
network = rx.generators.erdos_renyi_gnp_random_graph(20, 0.3, seed=42)

# Calculate centrality for coloring
centrality = rx.betweenness_centrality(network)
centrality_values = [centrality[node] for node in network.node_indices()]

# Use colormap for nodes
plt.figure(figsize=(10, 8))
rx.visualization.mpl_draw(
    network,
    pos=rx.spring_layout(network, seed=42),
    node_color=centrality_values,
    cmap='viridis',
    node_size=500,
    with_labels=True,
    font_size=8,
    font_color='white'
)

# Add colorbar
plt.colorbar(plt.cm.ScalarMappable(cmap='viridis'), label='Betweenness Centrality')
plt.title("Network Colored by Betweenness Centrality")
plt.axis('off')
plt.show()

# Compare different graph structures
graphs = {
    'Complete': rx.generators.complete_graph(6),
    'Cycle': rx.generators.cycle_graph(6),
    'Star': rx.generators.star_graph(6),
    'Path': rx.generators.path_graph(6)
}

fig, axes = plt.subplots(2, 2, figsize=(12, 12))
axes = axes.flatten()

for i, (name, graph) in enumerate(graphs.items()):
    pos = rx.circular_layout(graph) if name != 'Star' else rx.spring_layout(graph)
    
    rx.visualization.mpl_draw(
        graph,
        pos=pos,
        ax=axes[i],
        node_color='lightblue',
        node_size=300,
        with_labels=True
    )
    axes[i].set_title(f"{name} Graph ({graph.num_nodes()} nodes, {graph.num_edges()} edges)")
    axes[i].axis('off')

plt.tight_layout()
plt.show()

# Professional visualization with Graphviz
hierarchical_graph = rx.PyDiGraph()
org_nodes = hierarchical_graph.add_nodes_from([
    'CEO', 'CTO', 'CFO', 'Dev1', 'Dev2', 'Acc1', 'Acc2'
])
org_edges = [
    (0, 1, 'reports_to'),  # CEO -> CTO
    (0, 2, 'reports_to'),  # CEO -> CFO  
    (1, 3, 'manages'),     # CTO -> Dev1
    (1, 4, 'manages'),     # CTO -> Dev2
    (2, 5, 'manages'),     # CFO -> Acc1
    (2, 6, 'manages'),     # CFO -> Acc2
]
hierarchical_graph.add_edges_from(org_edges)

# Node styling function
def node_attrs(node_data):
    role = ['CEO', 'CTO', 'CFO', 'Dev1', 'Dev2', 'Acc1', 'Acc2'][node_data]
    if role == 'CEO':
        return {'shape': 'box', 'style': 'filled', 'fillcolor': 'gold', 'fontweight': 'bold'}
    elif role in ['CTO', 'CFO']:
        return {'shape': 'box', 'style': 'filled', 'fillcolor': 'lightblue'}
    else:
        return {'shape': 'ellipse', 'style': 'filled', 'fillcolor': 'lightgreen'}

# Edge styling function  
def edge_attrs(edge_data):
    if edge_data == 'reports_to':
        return {'color': 'red', 'style': 'bold'}
    else:
        return {'color': 'blue'}

# Generate DOT source
dot_source = rx.visualization.graphviz_draw(
    hierarchical_graph,
    node_attr_fn=lambda node: node_attrs(node),
    edge_attr_fn=lambda edge: edge_attrs(edge),
    graph_attr={'rankdir': 'TB', 'splines': 'ortho'},  # Top-to-bottom, orthogonal edges
    method='dot'  # Hierarchical layout
)

print("Generated DOT source:")
print(dot_source)

# Save to file (requires graphviz installation)
# rx.visualization.graphviz_draw(
#     hierarchical_graph,
#     node_attr_fn=node_attrs,
#     edge_attr_fn=edge_attrs,
#     filename='org_chart.png',
#     image_type='png',
#     method='dot'
# )

# Create interactive plot with hover information
network = rx.generators.barabasi_albert_graph(30, 2, seed=42)
pos = rx.spring_layout(network, seed=42)

# Calculate node properties for display
degrees = {node: network.degree(node) for node in network.node_indices()}
max_degree = max(degrees.values())

# Size nodes by degree
node_sizes = [300 + (degrees[node] / max_degree) * 700 for node in network.node_indices()]

plt.figure(figsize=(12, 10))
scatter = plt.scatter(
    [pos[node][0] for node in network.node_indices()],
    [pos[node][1] for node in network.node_indices()],
    s=node_sizes,
    c=[degrees[node] for node in network.node_indices()],
    cmap='plasma',
    alpha=0.7
)

# Draw edges
for source, target in network.edge_list():
    x1, y1 = pos[source]
    x2, y2 = pos[target]
    plt.plot([x1, x2], [y1, y2], 'k-', alpha=0.3, linewidth=0.5)

# Add colorbar and labels
plt.colorbar(scatter, label='Node Degree')
plt.title("Scale-Free Network (Node size and color = degree)")
plt.axis('off')

# Add node labels for high-degree nodes
for node in network.node_indices():
    if degrees[node] > max_degree * 0.7:  # Label high-degree nodes
        x, y = pos[node]
        plt.annotate(f'{node}\n(deg={degrees[node]})', 
                    (x, y), xytext=(5, 5), textcoords='offset points',
                    fontsize=8, bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8))

plt.show()

import matplotlib.animation as animation

# Create sequence of graphs showing growth
def create_growth_sequence(final_size=20, steps=10):
    """Create sequence showing preferential attachment growth."""
    sequences = []
    
    for i in range(1, steps + 1):
        size = max(3, int(final_size * i / steps))
        graph = rx.generators.barabasi_albert_graph(size, 2, seed=42)
        sequences.append(graph)
    
    return sequences

# Generate growth sequence
growth_graphs = create_growth_sequence(20, 10)

# Create animation
fig, ax = plt.subplots(figsize=(10, 8))

def animate(frame):
    ax.clear()
    graph = growth_graphs[frame]
    pos = rx.spring_layout(graph, seed=42)
    
    rx.visualization.mpl_draw(
        graph,
        pos=pos,
        ax=ax,
        node_color='lightblue',
        node_size=300,
        with_labels=True,
        font_size=8
    )
    
    ax.set_title(f"Preferential Attachment Growth: {graph.num_nodes()} nodes, {graph.num_edges()} edges")
    ax.axis('off')

# Create animation
anim = animation.FuncAnimation(fig, animate, frames=len(growth_graphs), 
                             interval=800, repeat=True)

plt.show()

# Save animation (optional)
# anim.save('graph_growth.gif', writer='pillow', fps=1)

def create_styled_visualization(graph, title="Styled Graph"):
    """Create a professionally styled graph visualization."""
    
    # Calculate layout
    pos = rx.spring_layout(graph, seed=42)
    
    # Calculate node properties
    degrees = {node: graph.degree(node) for node in graph.node_indices()}
    centrality = rx.betweenness_centrality(graph)
    
    # Style configuration
    node_sizes = [200 + degrees[node] * 100 for node in graph.node_indices()]
    node_colors = [centrality[node] for node in graph.node_indices()]
    
    # Create visualization
    plt.figure(figsize=(12, 10))
    
    # Draw edges first (behind nodes)
    for source, target in graph.edge_list():
        x1, y1 = pos[source]
        x2, y2 = pos[target]
        plt.plot([x1, x2], [y1, y2], 'gray', alpha=0.4, linewidth=1)
    
    # Draw nodes with custom styling
    scatter = plt.scatter(
        [pos[node][0] for node in graph.node_indices()],
        [pos[node][1] for node in graph.node_indices()],
        s=node_sizes,
        c=node_colors,
        cmap='coolwarm',
        alpha=0.8,
        edgecolors='black',
        linewidths=1
    )
    
    # Add labels for important nodes
    important_nodes = [node for node in graph.node_indices() 
                      if centrality[node] > 0.1]
    
    for node in important_nodes:
        x, y = pos[node]
        plt.annotate(f'{node}', (x, y), ha='center', va='center',
                    fontweight='bold', fontsize=10,
                    bbox=dict(boxstyle='circle,pad=0.3', facecolor='white', alpha=0.8))
    
    # Styling
    plt.colorbar(scatter, label='Betweenness Centrality', shrink=0.8)
    plt.title(title, fontsize=16, fontweight='bold', pad=20)
    plt.axis('off')
    
    # Add legend
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='gray', 
                  markersize=8, label='Small degree'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='gray',
                  markersize=12, label='Large degree')
    ]
    plt.legend(handles=legend_elements, loc='upper right')
    
    plt.tight_layout()
    return plt.gcf()

# Apply to different graph types
test_graphs = [
    (rx.generators.karate_club_graph(), "Karate Club Network"),
    (rx.generators.barabasi_albert_graph(30, 2), "Scale-Free Network"),
    (rx.generators.erdos_renyi_gnp_random_graph(25, 0.15), "Random Network")
]

for graph, title in test_graphs:
    fig = create_styled_visualization(graph, title)
    plt.show()