Update my_pages/multiverse.py

my_pages/multiverse.py

@@ -2,67 +2,93 @@ import streamlit as st
 import plotly.graph_objects as go
 from utils import development_stages
 
-def render():
-    st.title("Multiverse of Developer Decisions")
+def build_tree_edges(selected_path=None):
+    edges = []
+    node_labels = ["Start"]
+    node_positions = [(0, 0)]
+    node_stage = [0]  # stage index for each node
+    node_id_map = {("Start", 0): 0}
+
+    # Keep track of path highlighting
+    highlight_edges = set()
+
+    prev_nodes = [0]  # Start node index
+
+    for stage_idx, stage in enumerate(development_stages, start=1):
+        next_nodes = []
+        options = stage["questions"]
+
+        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
+
+                edges.append((parent_idx, node_index))
+
+                # 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))
 
-    # Extract stage labels
-    stage_labels = [stage["label"] for stage in development_stages]
+                next_nodes.append(node_index)
 
-    # Generate options from each stage's questions
-    stage_options = {stage["label"]: stage["questions"] for stage in development_stages}
+        prev_nodes = [n for n in next_nodes if n not in prev_nodes]
 
-    # Let user pick one option from each stage
+    return node_labels, node_positions, edges, highlight_edges
+
+def render():
+    st.title("Multiverse of Developer Decisions - Tree View")
+
+    # User selection for each stage
     selected_path = []
-    st.subheader("Choose your path")
-    for stage_label, options in stage_options.items():
-        choice = st.selectbox(f"{stage_label}:", options, key=stage_label)
+    for stage in development_stages:
+        choice = st.selectbox(f"{stage['icon']} {stage['label']}", stage["questions"], key=stage["label"])
         selected_path.append(choice)
 
-    # Prepare Sankey diagram labels
-    labels = []
-    for stage_label in stage_labels:
-        labels.extend(stage_options[stage_label])
-
-    # Create links between options in consecutive stages
-    source = []
-    target = []
-    value = []
-    colors = []
-
-    options_per_stage = [stage_options[stage] for stage in stage_labels]
-
-    for i in range(len(stage_labels) - 1):
-        start_idx = sum(len(opts) for opts in options_per_stage[:i])
-        next_start_idx = sum(len(opts) for opts in options_per_stage[:i+1])
-
-        for s in range(start_idx, start_idx + len(options_per_stage[i])):
-            for t in range(next_start_idx, next_start_idx + len(options_per_stage[i+1])):
-                source.append(s)
-                target.append(t)
-                value.append(1)  # constant thickness for now
-
-                # Highlight chosen path
-                if labels[s] == selected_path[i] and labels[t] == selected_path[i+1]:
-                    colors.append("rgba(0, 150, 0, 0.8)")
-                else:
-                    colors.append("rgba(200, 200, 200, 0.3)")
-
-    # Create Sankey diagram
-    fig = go.Figure(data=[go.Sankey(
-        arrangement="snap",
-        node=dict(
-            pad=15,
-            thickness=20,
-            line=dict(color="black", width=0.5),
-            label=labels,
-            color="rgba(100, 100, 200, 0.8)"
-        ),
-        link=dict(
-            source=source,
-            target=target,
-            value=value,
-            color=colors
-        )
-    )])
+    labels, positions, edges, highlight_edges = build_tree_edges(selected_path)
+
+    # Draw nodes
+    x_vals = [pos[0] for pos in positions]
+    y_vals = [pos[1] for pos in positions]
+
+    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'))
+    )
+
+    # Draw edges with highlighting
+    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
+        edge_traces.append(go.Scatter(
+            x=[positions[src][0], positions[dst][0]],
+            y=[positions[src][1], positions[dst][1]],
+            mode='lines',
+            line=dict(width=width, color=color),
+            hoverinfo='none'
+        ))
+
+    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)
+    )
 
     st.plotly_chart(fig, use_container_width=True)