Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import time
 import plotly.graph_objects as go
+import random
 
 # Set page layout to wide mode
 st.set_page_config(layout="wide")
@@ -54,29 +55,11 @@ possibilities = {
     ]
 }
 
-# Initialize session state for patient counts if not already present
-if 'well_count' not in st.session_state:
-    st.session_state['well_count'] = 0
-if 'slight_improvement_count' not in st.session_state:
-    st.session_state['slight_improvement_count'] = 0
-if 'no_response_count' not in st.session_state:
-    st.session_state['no_response_count'] = 0
-
-# Function to simulate patient response
-def simulate_patient_response():
-    # Simulating possible outcomes - can be adjusted for more complexity
-    # Example random patient responses (you can adjust this logic)
-    import random
-    outcome = random.choice(["well", "slight improvement", "no response"])
-    if outcome == "well":
-        st.session_state['well_count'] += 1
-    elif outcome == "slight improvement":
-        st.session_state['slight_improvement_count'] += 1
-    elif outcome == "no response":
-        st.session_state['no_response_count'] += 1
-
-# Create a tree diagram using Plotly (unchanged from original)
-def create_tree_diagram(stage):
+# Initialize a dictionary to keep track of patient counts for each possibility
+patient_counts = {i: [0] * len(possibilities[i]) for i in possibilities}
+
+# Create a tree diagram using Plotly
+def create_tree_diagram(stage, selected_box=None):
     labels = [f"Stage {stage}: {outcomes[stage - 1][0]}"]
     parents = [""]
     values = [1]  # Root node value
@@ -84,18 +67,18 @@ def create_tree_diagram(stage):
     
     stage_possibilities = possibilities.get(stage, [("No specific possibilities defined", "info")])
     
-    for possibility, outcome_type in stage_possibilities:
-        labels.append(possibility)
+    for idx, (possibility, outcome_type) in enumerate(stage_possibilities):
+        labels.append(f"{possibility} - Patients: {patient_counts[stage][idx]}")
         parents.append(f"Stage {stage}: {outcomes[stage - 1][0]}")
         values.append(1)  # Equal weight for all possibilities
         if outcome_type == "success":
-            colors.append('#d4edda')
+            colors.append('#d4edda' if idx != selected_box else '#28a745')
         elif outcome_type == "info":
-            colors.append('#cce5ff')
+            colors.append('#cce5ff' if idx != selected_box else '#007bff')
         elif outcome_type == "warning":
-            colors.append('#fff3cd')
+            colors.append('#fff3cd' if idx != selected_box else '#ffc107')
         elif outcome_type == "error":
-            colors.append('#f8d7da')
+            colors.append('#f8d7da' if idx != selected_box else '#dc3545')
         else:
             colors.append('lightgrey')
     
@@ -116,48 +99,61 @@ def create_tree_diagram(stage):
     
     return fig
 
+# Function to increment patient counts in random boxes for each stage
+def run_simulation():
+    for stage in possibilities:
+        selected_box = random.randint(0, len(possibilities[stage]) - 1)
+        patient_counts[stage][selected_box] += 1
+        st.session_state[f'selected_box_{stage}'] = selected_box
+
+# Initialize the selected box in session state if it doesn't exist
+for stage in possibilities:
+    if f'selected_box_{stage}' not in st.session_state:
+        st.session_state[f'selected_box_{stage}'] = None
+
 # Streamlit app layout
 st.title("Precision Medicine AI Agents - Treatment Decision Tree")
 
-# Display the current patient counts
-st.write(f"### Well: {st.session_state['well_count']}")
-st.write(f"### Slight Improvement: {st.session_state['slight_improvement_count']}")
-st.write(f"### No Response: {st.session_state['no_response_count']}")
-
-# Add a start animation button
-start_button = st.button("Start Animation")
-
-if start_button:
-    simulate_patient_response()  # Simulate patient response when button is pressed
-
-    for i, (outcome, outcome_type) in enumerate(outcomes, 1):
-        with st.container():
-            col1, col2 = st.columns([1, 2])
-
-            with col1:
-                if outcome_type == "success":
-                    st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#d4edda; color:#155724;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
-                elif outcome_type == "info":
-                    st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#cce5ff; color:#004085;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
-                elif outcome_type == "warning":
-                    st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#fff3cd; color:#856404;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
-                elif outcome_type == "error":
-                    st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#f8d7da; color:#721c24;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
-
-                st.markdown("### Possibilities:")
-                stage_possibilities = possibilities.get(i, [("No specific possibilities defined", "info")])
-                for possibility, possibility_type in stage_possibilities:
-                    if possibility_type == "success":
-                        st.markdown(f"<div style='padding:5px; background-color:#d4edda; color:#155724;'><strong>{possibility}</strong></div>", unsafe_allow_html=True)
-                    elif possibility_type == "info":
-                        st.markdown(f"<div style='padding:5px; background-color:#cce5ff; color:#004085;'><strong>{possibility}</strong></div>", unsafe_allow_html=True)
-                    elif possibility_type == "warning":
-                        st.markdown(f"<div style='padding:5px; background-color:#fff3cd; color:#856404;'><strong>{possibility}</strong></div>", unsafe_allow_html=True)
-                    elif possibility_type == "error":
-                        st.markdown(f"<div style='padding:5px; background-color:#f8d7da; color:#721c24;'><strong>{possibility}</strong></div>", unsafe_allow_html=True)
-
-            with col2:
-                fig = create_tree_diagram(i)
-                st.plotly_chart(fig, use_container_width=True)
-
-        time.sleep(5)
+# Add a "Run Simulation" button
+run_button = st.button("Run Simulation")
+
+# If the button is pressed, run the simulation
+if run_button:
+    run_simulation()
+
+# Render the animation with the updated patient counts
+for i, (outcome, outcome_type) in enumerate(outcomes, 1):
+    with st.container():
+        col1, col2 = st.columns([1, 2])
+
+        with col1:
+            if outcome_type == "success":
+                st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#d4edda; color:#155724;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
+            elif outcome_type == "info":
+                st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#cce5ff; color:#004085;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
+            elif outcome_type == "warning":
+                st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#fff3cd; color:#856404;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
+            elif outcome_type == "error":
+                st.markdown(f"<div style='padding:10px; border-radius:5px; background-color:#f8d7da; color:#721c24;'><strong>Stage {i}:</strong> {outcome}</div>", unsafe_allow_html=True)
+
+            st.markdown("### Possibilities:")
+            stage_possibilities = possibilities.get(i, [("No specific possibilities defined", "info")])
+            for idx, (possibility, possibility_type) in enumerate(stage_possibilities):
+                color = 'lightgrey'
+                if possibility_type == "success":
+                    color = '#d4edda'
+                elif possibility_type == "info":
+                    color = '#cce5ff'
+                elif possibility_type == "warning":
+                    color = '#fff3cd'
+                elif possibility_type == "error":
+                    color = '#f8d7da'
+
+                # Display the possibility and its patient count
+                st.markdown(f"<div style='padding:5px; background-color:{color};'><strong>{possibility} - Patients: {patient_counts[i][idx]}</strong></div>", unsafe_allow_html=True)
+
+        with col2:
+            selected_box = st.session_state[f'selected_box_{i}']
+            fig = create_tree_diagram(i, selected_box)
+            st.plotly_chart(fig, use_container_width=True)
+