Update my_pages/multiverse.py

my_pages/multiverse.py

@@ -1,77 +1,115 @@
+# pages/multiverse.py
 import streamlit as st
 import plotly.graph_objects as go
 from utils import development_stages
 
-def build_tree_edges_and_path(selected_path):
-    edges = []
+def build_tree_and_trace_path(selected_path):
+    """
+    Build tree nodes and edges. Then trace selected_path (one choice per stage)
+    by walking children of the current node to find the matching label at each stage.
+    Returns: node_labels, node_positions, edges, highlight_edges, highlight_nodes
+    """
     node_labels = ["Start"]
-    node_positions = [(0, 0)]
+    node_positions = [(0.0, 0.0)]
     node_stage = [0]
-    node_index_map = {(0, "Start"): 0}  # (stage_idx, label) → index
+    edges = []
 
-    prev_nodes = [0]  # start node
-    stage_nodes_map = {}  # stage_idx → [node indices]
+    prev_nodes = [0]  # nodes at previous stage (start)
+    y_spacing_base = 1.0
 
-    # Build all nodes and edges
+    # Build nodes and edges stage by stage
     for stage_idx, stage in enumerate(development_stages, start=1):
         options = stage["questions"]
-        stage_nodes_map[stage_idx] = []
+        next_nodes = []
+
+        # For layout niceness: compute a span for children under each parent
         for parent_idx in prev_nodes:
             px, py = node_positions[parent_idx]
+            # Spread children vertically around the parent's y
+            span = max(1, len(options))
             for opt_idx, opt in enumerate(options):
-                child_x = stage_idx
-                child_y = py + (opt_idx - len(options) / 2) * (1.5 / (stage_idx))
+                child_x = float(stage_idx)
+                # center children around px, offset by opt_idx
+                offset = (opt_idx - (len(options) - 1) / 2.0) * (y_spacing_base / (stage_idx))
+                child_y = py + offset
                 node_index = len(node_labels)
                 node_labels.append(opt)
                 node_positions.append((child_x, child_y))
                 node_stage.append(stage_idx)
-                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]
+                next_nodes.append(node_index)
 
-    # Determine the exact path
+        # After all parents, the next stage's parents are all next_nodes
+        prev_nodes = next_nodes
+
+    # Trace the single chosen path by walking children
     highlight_edges = set()
-    highlight_nodes = set([0])  # start node always highlighted
+    highlight_nodes = set([0])
     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
+
+    for stage_idx, chosen_label in enumerate(selected_path, start=1):
+        # children of current_node
+        children = [dst for (src, dst) in edges if src == current_node]
+        # find a child among these whose label equals chosen_label
+        found_child = None
+        for c in children:
+            if node_labels[c] == chosen_label:
+                found_child = c
+                break
+        if found_child is None:
+            # chosen label not found under this parent — stop tracing
+            break
+        highlight_edges.add((current_node, found_child))
+        highlight_nodes.add(found_child)
+        current_node = found_child
 
     return node_labels, node_positions, edges, highlight_edges, highlight_nodes
 
+
 def render():
-    st.title("Multiverse of Developer Decisions - Tree View")
+    st.title("Multiverse of Developer Decisions — Tree View")
 
-    # User selection for each stage
+    # --- User picks one choice per stage via dropdowns ---
     selected_path = []
     for stage in development_stages:
-        choice = st.selectbox(f"{stage['icon']} {stage['label']}", stage["questions"], key=stage["label"])
+        # use a stable key per stage to avoid conflicts
+        key = f"multiverse_choice_{stage['label']}"
+        choice = st.selectbox(f"{stage['icon']} {stage['label']}", stage["questions"], key=key)
         selected_path.append(choice)
 
-    labels, positions, edges, highlight_edges, highlight_nodes = build_tree_edges_and_path(selected_path)
+    # --- Build tree and compute which edges/nodes to highlight ---
+    labels, positions, edges, highlight_edges, highlight_nodes = build_tree_and_trace_path(selected_path)
 
+    # --- Prepare node and edge traces for Plotly ---
     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_colors = []
+    for idx in range(len(labels)):
+        if idx in highlight_nodes:
+            node_colors.append("rgba(34,139,34,0.95)")  # green for selected path nodes
+        elif idx == 0:
+            node_colors.append("rgba(30,144,255,0.9)")  # start node distinct
+        else:
+            node_colors.append("rgba(135,206,250,0.6)")  # default skyblue
 
     node_trace = go.Scatter(
         x=x_vals, y=y_vals,
         mode='markers+text',
         text=labels,
         textposition="top center",
-        marker=dict(size=20, color=node_colors, line=dict(width=2, color='black'))
+        marker=dict(size=18, color=node_colors, line=dict(width=1, color='black')),
+        hoverinfo="text"
     )
 
     edge_traces = []
     for src, dst in edges:
-        color = 'green' if (src, dst) in highlight_edges else 'gray'
-        width = 3 if (src, dst) in highlight_edges else 1.5
+        if (src, dst) in highlight_edges:
+            color = "rgba(0,128,0,0.9)"  # bright green
+            width = 4
+        else:
+            color = "rgba(120,120,120,0.4)"
+            width = 1.5
         edge_traces.append(go.Scatter(
             x=[positions[src][0], positions[dst][0]],
             y=[positions[src][1], positions[dst][1]],
@@ -80,13 +118,15 @@ def render():
             hoverinfo='none'
         ))
 
+    # --- Render figure ---
     fig = go.Figure(data=edge_traces + [node_trace])
     fig.update_layout(
         showlegend=False,
         xaxis=dict(visible=False),
         yaxis=dict(visible=False),
         plot_bgcolor='white',
-        margin=dict(l=0, r=0, t=0, b=0)
+        margin=dict(l=10, r=10, t=10, b=10),
+        hovermode="closest"
     )
 
     st.plotly_chart(fig, use_container_width=True)