Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import time
-import random
 import plotly.graph_objects as go
 
 # Set page layout to wide mode
@@ -20,114 +19,145 @@ outcomes = [
 
 # Different possible outcomes of each treatment
 possibilities = {
-    1: ["Well", "Slight Improvement", "No Response"],
-    2: ["Significant Improvement", "Mild Side Effects"],
-    3: ["Effective Treatment", "Inconclusive Results"],
-    4: ["Improvement", "Significant Side Effects"],
-    5: ["Worsened Side Effects", "Manageable Side Effects"],
-    6: ["Successful Adjustment", "Further Adjustment Needed"],
-    7: ["Good Response", "Limited Response"],
-    8: ["Complete Recovery", "Partial Improvement"]
+    1: [
+        ("Possibility 1: Patient responds well to Treatment 1", "success"),
+        ("Possibility 2: Slight improvement, but inconclusive results", "info"),
+        ("Possibility 3: No response, reevaluation needed", "error")
+    ],
+    2: [
+        ("Possibility 1: Significant improvement", "success"),
+        ("Possibility 2: Mild side effects", "warning")
+    ],
+    3: [
+        ("Possibility 1: Treatment is effective", "success"),
+        ("Possibility 2: Inconclusive lab results", "info")
+    ],
+    4: [
+        ("Possibility 1: Improvement with Drug B", "success"),
+        ("Possibility 2: Significant side effects", "error")
+    ],
+    5: [
+        ("Possibility 1: Side effects worsen, modify dosage", "warning"),
+        ("Possibility 2: Manageable side effects", "info")
+    ],
+    6: [
+        ("Possibility 1: Dosage adjustment successful", "success"),
+        ("Possibility 2: Further modification needed", "warning")
+    ],
+    7: [
+        ("Possibility 1: Patient responds well to modified treatment", "success"),
+        ("Possibility 2: Limited response, consider alternatives", "warning")
+    ],
+    8: [
+        ("Possibility 1: Complete recovery", "success"),
+        ("Possibility 2: Partial improvement, continue monitoring", "info")
+    ]
 }
 
-# Store patients' records (persistent across runs)
-patients_records = {
-    "Well": 0,
-    "Slight Improvement": 0,
-    "No Response": 0,
-    "Significant Improvement": 0,
-    "Mild Side Effects": 0,
-    "Effective Treatment": 0,
-    "Inconclusive Results": 0,
-    "Improvement": 0,
-    "Significant Side Effects": 0,
-    "Worsened Side Effects": 0,
-    "Manageable Side Effects": 0,
-    "Successful Adjustment": 0,
-    "Further Adjustment Needed": 0,
-    "Good Response": 0,
-    "Limited Response": 0,
-    "Complete Recovery": 0,
-    "Partial Improvement": 0
-}
-
-# Generate random patient names
-def generate_patient_name():
-    first_names = ["John", "Jane", "Alex", "Sara", "Tom", "Emily", "Michael", "Lucy"]
-    last_names = ["Smith", "Doe", "Brown", "Wilson", "Taylor", "Lee"]
-    return f"{random.choice(first_names)} {random.choice(last_names)}"
-
-# Run simulation and update patient record based on random outcomes
-def simulate_patient_outcome():
-    for stage in range(1, 9):
-        patient_name = generate_patient_name()
-        outcome_list = possibilities[stage]
-        outcome = random.choice(outcome_list)
-        patients_records[outcome] += 1  # Increment count for that outcome
-
-# Create a tree diagram using Plotly with patient information
-def create_tree_diagram():
-    labels = ["Patients"]
+# 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):
+    labels = [f"Stage {stage}: {outcomes[stage - 1][0]}"]
     parents = [""]
-    values = [sum(patients_records.values())]
-    colors = ['lightgrey']
+    values = [1]  # Root node value
+    colors = ['lightgrey']  # Root node color
     
-    for category, patient_count in patients_records.items():
-        labels.append(f"{category} ({patient_count} patients)")
-        parents.append("Patients")
-        values.append(patient_count)
-        colors.append('lightgrey')
+    stage_possibilities = possibilities.get(stage, [("No specific possibilities defined", "info")])
+    
+    for possibility, outcome_type in stage_possibilities:
+        labels.append(possibility)
+        parents.append(f"Stage {stage}: {outcomes[stage - 1][0]}")
+        values.append(1)  # Equal weight for all possibilities
+        if outcome_type == "success":
+            colors.append('#d4edda')
+        elif outcome_type == "info":
+            colors.append('#cce5ff')
+        elif outcome_type == "warning":
+            colors.append('#fff3cd')
+        elif outcome_type == "error":
+            colors.append('#f8d7da')
+        else:
+            colors.append('lightgrey')
     
     fig = go.Figure(go.Treemap(
         labels=labels,
         parents=parents,
         values=values,
         marker=dict(colors=colors),
-        textinfo="label+value",
-        hoverinfo="label+value"
+        textinfo="label",
+        textfont=dict(size=14),
     ))
 
     fig.update_layout(
         margin=dict(t=10, l=10, r=10, b=10),
         width=600, height=400,
-        uniformtext=dict(minsize=12, mode='show'),
-        transition_duration=500
+        uniformtext=dict(minsize=12, mode='show')
     )
     
     return fig
 
 # Streamlit app layout
-st.title("Precision Medicine AI Agents - Treatment Decision Tree Simulation")
+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
-run_simulation = st.button("Run Simulation")
-
-if run_simulation:
-    simulate_patient_outcome()  # Simulate new patient outcomes
-
-# Display the patient outcomes incrementally
-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("### Current Patients:")
-            stage_possibilities = possibilities[i]
-            for possibility in stage_possibilities:
-                patient_count = patients_records.get(possibility, 0)
-                st.markdown(f"**{possibility}:** {patient_count} patients")
-
-        with col2:
-            fig = create_tree_diagram()  # Update chart with patient records
-            st.plotly_chart(fig, use_container_width=True)
-
-        time.sleep(1)  # Shorter delay for smoother transitions
+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)