Update app.py

app.py

@@ -1,137 +1,157 @@
 import streamlit as st
-import time
-import plotly.graph_objects as go
-
-# Set page layout to wide mode
-st.set_page_config(layout="wide")
-
-# Treatment outcomes and possibilities
-outcomes = [
-    ("Doctor prescribes Treatment 1 (Personalized Drug A)", "info"),
-    ("Nurse monitors Patient - Slight improvement", "success"),
-    ("Clinician evaluates lab results - Treatment seems effective", "info"),
-    ("Doctor prescribes Treatment 2 (Personalized Drug B)", "warning"),
-    ("Nurse monitors Patient - Side effects observed", "error"),
-    ("Clinician suggests modifying dosage", "warning"),
-    ("Doctor prescribes modified Treatment 2", "info"),
-    ("Patient shows significant improvement", "success")
-]
-
-# Different possible outcomes of each treatment
-possibilities = {
-    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")
-    ]
+import numpy as np
+import random
+
+# Initialize constants
+DOCTOR_ACTIONS = ["Prescribe Medication", "Recommend Tests", "Consult Clinician", "Schedule Surgery"]
+NURSE_ACTIONS = ["Monitor Vitals", "Administer Medication", "Report to Doctor", "Assist Surgery"]
+PATIENT_CONDITIONS = ["Healthy", "Mild Illness", "Chronic Illness", "Emergency"]
+DOCTOR_EMOTIONS = ["Calm", "Stressed", "Overwhelmed"]
+NURSE_EMOTIONS = ["Focused", "Fatigued", "Panicked"]
+
+# Rewards and Penalties
+REWARDS = {
+    "Prescribe Medication": 12,
+    "Recommend Tests": 7,
+    "Consult Clinician": 9,
+    "Schedule Surgery": 17,
+    "Monitor Vitals": 4,
+    "Administer Medication": 12,
 }
+PENALTIES = {
+    "Wrong Medication": -5,
+    "Missed Diagnosis": -10,
+    "Stress-Induced Mistake": -7,
+}
+
+# Initialize session state for metrics and satisfaction
+if "performance_metrics" not in st.session_state:
+    st.session_state.performance_metrics = {
+        "Doctor": {"successful_treatments": 0, "failed_treatments": 0},
+        "Nurse": {"successful_assists": 0, "failed_assists": 0}
+    }
+if "patient_satisfaction" not in st.session_state:
+    st.session_state.patient_satisfaction = 100
 
-# Create a tree diagram using Plotly with animations
-def create_tree_diagram(stage):
-    labels = [f"Stage {stage}: {outcomes[stage - 1][0]}"]
-    parents = [""]  # Root node
-    values = [1]  # Root node value
-    colors = ['lightgrey']  # Root node color
+
+# Define the Agent class for Q-learning
+class Agent:
+    def __init__(self, agent_type):
+        self.agent_type = agent_type
+        self.actions = DOCTOR_ACTIONS if agent_type == "Doctor" else NURSE_ACTIONS
+        self.q_table = np.zeros((len(PATIENT_CONDITIONS), len(self.actions)))
     
-    stage_possibilities = possibilities.get(stage, [("No specific possibilities defined", "info")])
+    def choose_action(self, state, exploration_rate=0.05):
+        if random.uniform(0, 1) < exploration_rate:  # Explore
+            return random.randint(0, len(self.actions)-1)
+        else:  # Exploit (choose best action)
+            return np.argmax(self.q_table[state])
+
+    def update_q_value(self, state, action, reward, learning_rate=0.1, discount_factor=0.9):
+        old_q_value = self.q_table[state, action]
+        best_future_q_value = np.max(self.q_table)
+        new_q_value = old_q_value + learning_rate * (reward + discount_factor * best_future_q_value - old_q_value)
+        self.q_table[state, action] = new_q_value
+
+
+# Instantiate agents
+doctor_agent = Agent("Doctor")
+nurse_agent = Agent("Nurse")
+
+
+# Function to simulate a special event
+def simulate_special_event():
+    event = random.choice([None, "Disease Outbreak", "Resource Shortage"])
+    if event == "Disease Outbreak":
+        st.subheader("Special Event: Disease Outbreak")
+        st.write("A sudden disease outbreak has flooded the hospital with new patients. Resources are limited!")
+    elif event == "Resource Shortage":
+        st.subheader("Special Event: Resource Shortage")
+        st.write("A medical supply shortage is impacting the hospital. Staff must prioritize high-risk patients.")
+    return event
+
+
+# Function to handle complications during treatment
+def handle_complications():
+    complication = random.choices(
+        [None, "Allergic Reaction", "Unexpected Complication"],
+        weights=[0.6, 0.2, 0.2]
+    )[0]
     
-    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')
+    penalty = 0
+    if complication:
+        st.subheader(f"Complication: {complication}")
+        if complication == "Allergic Reaction":
+            st.write("The patient has developed an allergic reaction to the prescribed medication!")
+            penalty = PENALTIES["Wrong Medication"]
+            st.session_state.patient_satisfaction -= 10
+        elif complication == "Unexpected Complication":
+            st.write("An unexpected complication occurred during surgery!")
+            penalty = PENALTIES["Stress-Induced Mistake"]
+            st.session_state.patient_satisfaction -= 15
+    return penalty
+
+
+# Main simulation button logic
+if st.button("Run Simulation"):
+    # Simulate a special event
+    special_event = simulate_special_event()
+    
+    # Patient condition simulation
+    patient_state = random.choice(PATIENT_CONDITIONS)
+    st.write(f"Simulated Patient Condition: {patient_state}")
+    patient_index = PATIENT_CONDITIONS.index(patient_state)
+    
+    # Doctor and nurse emotions
+    doctor_emotion = random.choice(DOCTOR_EMOTIONS)
+    nurse_emotion = random.choice(NURSE_EMOTIONS)
+    st.write(f"Doctor's Emotional State: {doctor_emotion}")
+    st.write(f"Nurse's Emotional State: {nurse_emotion}")
+    
+    # Doctor Action
+    doctor_action_index = doctor_agent.choose_action(patient_index)
+    doctor_action = DOCTOR_ACTIONS[doctor_action_index]
+    st.write(f"Doctor's Chosen Action: {doctor_action}")
+    
+    # Nurse Action
+    nurse_action_index = nurse_agent.choose_action(patient_index)
+    nurse_action = NURSE_ACTIONS[nurse_action_index]
+    st.write(f"Nurse's Chosen Action: {nurse_action}")
+    
+    # Handle potential complications
+    penalty = handle_complications()
     
-    fig = go.Figure(go.Treemap(
-        labels=labels,
-        parents=parents,
-        values=values,
-        marker=dict(colors=colors),
-        textinfo="label",
-        textfont=dict(size=14),
-        hoverinfo="label+value+percent entry"
-    ))
-
-    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  # Animation speed
-    )
+    # Reward or penalty
+    reward = REWARDS.get(doctor_action, 0) if penalty == 0 else penalty
+    if doctor_emotion in ["Stressed", "Overwhelmed"]:
+        penalty += PENALTIES["Stress-Induced Mistake"]
     
-    return fig
-
-# Streamlit app layout
-st.title("Precision Medicine AI Agents - Treatment Decision Tree with Animations")
-
-# Add a start animation button
-start_button = st.button("Start Animation")
-
-if start_button:
-    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(3)  # Shorter delay for smoother transitions
+    # Update Q-values
+    doctor_agent.update_q_value(patient_index, doctor_action_index, reward)
+    
+    st.write(f"Doctor's Reward/Penalty: {reward}")
+    st.write(f"Patient Satisfaction: {st.session_state.patient_satisfaction}")
+    
+    # Outcome and Performance Metrics Update
+    outcome = random.choices(
+        ["Recovery", "Further Treatment Needed", "Complication"],
+        weights=[0.6, 0.3, 0.1]
+    )[0]
+    
+    st.write(f"Patient Status after Treatment: {outcome}")
+    
+    if outcome == "Recovery":
+        st.session_state.performance_metrics["Doctor"]["successful_treatments"] += 1
+        st.session_state.performance_metrics["Nurse"]["successful_assists"] += 1
+    else:
+        st.session_state.performance_metrics["Doctor"]["failed_treatments"] += 1
+        st.session_state.performance_metrics["Nurse"]["failed_assists"] += 1
+
+    st.write("Simulation completed! Run again for different outcomes.")
+
+# Display performance metrics
+st.subheader("Performance Metrics:")
+st.write(f"Doctor's Successful Treatments: {st.session_state.performance_metrics['Doctor']['successful_treatments']}")
+st.write(f"Doctor's Failed Treatments: {st.session_state.performance_metrics['Doctor']['failed_treatments']}")
+st.write(f"Nurse's Successful Assists: {st.session_state.performance_metrics['Nurse']['successful_assists']}")
+st.write(f"Nurse's Failed Assists: {st.session_state.performance_metrics['Nurse']['failed_assists']}")