visualization.py

import plotly.graph_objects as go
import networkx as nx
from typing import List, Set

# Updated color scheme per user request
COLOR_PATH = '#d62728'      # Red for path nodes & edges
COLOR_VISITED = '#2ca02c'   # Green for visited (non-path)
COLOR_UNTOUCHED = '#87cefa' # Light blue for untouched

def build_plot(G: nx.DiGraph, path: List[int], visited: List[int], start: int, intermediate, goals: Set[int], theme: str = 'Dark'):
    # Fixed layout for consistency (user-specified parameters)
    pos = nx.spring_layout(G, k=122, seed=42)

    path_set = set(path)
    visited_set = set(visited)
    path_edge_pairs = set()
    if len(path) > 1:
        path_edge_pairs = {(path[i], path[i+1]) for i in range(len(path)-1)}

    node_x = []
    node_y = []
    node_color = []
    node_text = []

    for node in G.nodes():
        x, y = pos[node]
        node_x.append(x)
        node_y.append(y)
        attrs = G.nodes[node]
        label = f"{node}: {attrs.get('name', node)}"
        if node in path_set:
            color = COLOR_PATH
        elif node in visited_set:
            color = COLOR_VISITED
        else:
            color = COLOR_UNTOUCHED
        node_color.append(color)
        node_text.append(label)

    # All edges (base layer)
    base_edge_x, base_edge_y = [], []
    path_edge_x, path_edge_y = [], []
    for u, v in G.edges():
        x0, y0 = pos[u]
        x1, y1 = pos[v]
        if (u, v) in path_edge_pairs:
            path_edge_x += [x0, x1, None]
            path_edge_y += [y0, y1, None]
        else:
            base_edge_x += [x0, x1, None]
            base_edge_y += [y0, y1, None]

    # Theme-dependent base styling
    if theme == 'Dark':
        base_edge_color = '#555555'
        marker_line_color = 'white'
        bg_color = '#111111'
        font_color = '#e0e0e0'
    else:
        base_edge_color = '#bbbbbb'
        marker_line_color = 'black'
        bg_color = '#ffffff'
        font_color = '#222222'

    base_edges = go.Scatter(
        x=base_edge_x, y=base_edge_y,
        line=dict(width=1, color=base_edge_color),
        hoverinfo='none', mode='lines', showlegend=False
    )
    path_edges = go.Scatter(
        x=path_edge_x, y=path_edge_y,
        line=dict(width=3, color=COLOR_PATH),
        hoverinfo='none', mode='lines', showlegend=False
    )

    node_trace = go.Scatter(
        x=node_x, y=node_y,
        mode='markers+text',
        text=node_text,
        textposition='top center',
        hoverinfo='text',
        marker=dict(color=node_color, size=28, line=dict(width=1.5, color=marker_line_color)),
        showlegend=False
    )

    # Prepare directional arrowheads using a separate scatter trace (triangle markers)
    arrow_x = []
    arrow_y = []
    arrow_color = []
    # Slight shorten factor to stop before node center
    shorten = 0.08
    for u, v in G.edges():
        x0, y0 = pos[u]
        x1, y1 = pos[v]
        dx = x1 - x0
        dy = y1 - y0
        adj_x1 = x1 - dx * shorten
        adj_y1 = y1 - dy * shorten
        arrow_x.append(adj_x1)
        arrow_y.append(adj_y1)
        arrow_color.append(COLOR_PATH if (u, v) in path_edge_pairs else '#888')

    arrow_trace = go.Scatter(
        x=arrow_x,
        y=arrow_y,
        mode='markers',
        hoverinfo='none',
        marker=dict(
            symbol='triangle-up',
            size=7,
            angle=0,
            color=arrow_color,
            line=dict(width=0)
        ),
        showlegend=False
    )
    # Build figure now with desired layering: base edges -> path edges -> arrowheads -> nodes
    fig = go.Figure()
    fig.add_trace(base_edges)
    fig.add_trace(path_edges)
    # fig.add_trace(arrow_trace)
    fig.add_trace(node_trace)
    # 3:2 aspect ratio & theming
    fig.update_layout(
        margin=dict(l=10, r=10, t=30, b=10),
        title="Pathway Search (3:2)",
        xaxis=dict(showgrid=False, zeroline=False, visible=False, scaleanchor='y', scaleratio=1.5),
        yaxis=dict(showgrid=False, zeroline=False, visible=False),
        paper_bgcolor=bg_color,
        plot_bgcolor=bg_color,
        font=dict(color=font_color)
    )
    return fig