Update my_pages/multiverse.py

my_pages/multiverse.py

@@ -2,49 +2,48 @@ import streamlit as st
 import plotly.graph_objects as go
 from utils import development_stages
 
-def build_tree_edges(selected_path=None):
+def build_tree_edges_and_path(selected_path):
     edges = []
     node_labels = ["Start"]
     node_positions = [(0, 0)]
-    node_stage = [0]  # stage index for each node
-    node_id_map = {("Start", 0): 0}
+    node_stage = [0]
+    node_index_map = {(0, "Start"): 0}  # (stage_idx, label) → index
 
-    # Keep track of path highlighting
-    highlight_edges = set()
-
-    prev_nodes = [0]  # Start node index
+    prev_nodes = [0]  # start node
+    stage_nodes_map = {}  # stage_idx → [node indices]
 
+    # Build all nodes and edges
     for stage_idx, stage in enumerate(development_stages, start=1):
-        next_nodes = []
         options = stage["questions"]
-
+        stage_nodes_map[stage_idx] = []
         for parent_idx in prev_nodes:
             px, py = node_positions[parent_idx]
             for opt_idx, opt in enumerate(options):
-                # Position children
                 child_x = stage_idx
                 child_y = py + (opt_idx - len(options) / 2) * (1.5 / (stage_idx))
                 node_index = len(node_labels)
                 node_labels.append(opt)
                 node_positions.append((child_x, child_y))
                 node_stage.append(stage_idx)
-                node_id_map[(opt, stage_idx)] = node_index
-
+                node_index_map[(stage_idx, opt)] = node_index
                 edges.append((parent_idx, node_index))
+                stage_nodes_map[stage_idx].append(node_index)
+        prev_nodes = stage_nodes_map[stage_idx]
 
-                # Highlight check
-                if selected_path and selected_path[stage_idx - 1] == opt and (
-                    (stage_idx == 1 and parent_idx == 0)
-                    or parent_idx in next_nodes
-                    or any(e[1] == parent_idx for e in edges)
-                ):
-                    highlight_edges.add((parent_idx, node_index))
-
-                next_nodes.append(node_index)
-
-        prev_nodes = [n for n in next_nodes if n not in prev_nodes]
-
-    return node_labels, node_positions, edges, highlight_edges
+    # Determine the exact path
+    highlight_edges = set()
+    highlight_nodes = set([0])  # start node always highlighted
+    current_node = 0
+    for stage_idx, choice in enumerate(selected_path, start=1):
+        if (stage_idx, choice) in node_index_map:
+            child_node = node_index_map[(stage_idx, choice)]
+            highlight_edges.add((current_node, child_node))
+            highlight_nodes.add(child_node)
+            current_node = child_node
+        else:
+            break  # choice not found — stop path
+
+    return node_labels, node_positions, edges, highlight_edges, highlight_nodes
 
 def render():
     st.title("Multiverse of Developer Decisions - Tree View")
@@ -55,21 +54,20 @@ def render():
         choice = st.selectbox(f"{stage['icon']} {stage['label']}", stage["questions"], key=stage["label"])
         selected_path.append(choice)
 
-    labels, positions, edges, highlight_edges = build_tree_edges(selected_path)
+    labels, positions, edges, highlight_edges, highlight_nodes = build_tree_edges_and_path(selected_path)
 
-    # Draw nodes
     x_vals = [pos[0] for pos in positions]
     y_vals = [pos[1] for pos in positions]
+    node_colors = ["green" if i in highlight_nodes else "skyblue" for i in range(len(labels))]
 
     node_trace = go.Scatter(
         x=x_vals, y=y_vals,
         mode='markers+text',
         text=labels,
         textposition="top center",
-        marker=dict(size=20, color='skyblue', line=dict(width=2, color='black'))
+        marker=dict(size=20, color=node_colors, line=dict(width=2, color='black'))
     )
 
-    # Draw edges with highlighting
     edge_traces = []
     for src, dst in edges:
         color = 'green' if (src, dst) in highlight_edges else 'gray'