zain329/EpidemicAI-Command-Center

1import streamlit as st
2import plotly.express as px
3import plotly.graph_objects as go
4import numpy as np
5import time
6
7# Import your custom backend
8from server.env import StratifiedEpidemicEnv, EpidemicAction
9from server.llm_agent import MultiAgentPolicySystem
10
11# --- 1. PAGE CONFIG & PREMIUM UX ---
12st.set_page_config(page_title="EpidemicAI Command Center", page_icon="🌍", layout="wide")
13
14# --- 2. SESSION STATE INITIALIZATION ---
15if "env" not in st.session_state:
16    st.session_state.env = StratifiedEpidemicEnv()
17if "agent" not in st.session_state:
18    st.session_state.agent = MultiAgentPolicySystem()
19if "history" not in st.session_state:
20    st.session_state.history = []
21if "scenario_text" not in st.session_state:
22    st.session_state.scenario_text = ""
23
24env = st.session_state.env
25agent = st.session_state.agent
26history = st.session_state.history
27
28# --- 3. THE SIDEBAR (CONTROLS & GOD MODE) ---
29st.sidebar.image("https://upload.wikimedia.org/wikipedia/commons/thumb/e/e5/Red_Cross_icon.svg/1024px-Red_Cross_icon.svg.png", width=50) # Optional logo
30st.sidebar.title("Simulation Controls")
31
32# Core Action
33if st.sidebar.button("▶️ Run Next Day", type="primary", use_container_width=True, disabled=env.current_day >= env.max_days):
34    with st.spinner("AI Cabinet is debating..."):
35        obs = env.state()
36        prev_act = env.prev_action if env.prev_action is not None else 0
37        
38        # 1. AI decides
39        action_dict = agent.get_action(obs, history, prev_act)
40        action_model = EpidemicAction(
41            reasoning=action_dict.get("reasoning", "Fallback logic"),
42            policy_choice=action_dict.get("policy_choice", 1)
43        )
44        
45        # 2. Environment steps
46        next_obs, reward, done, info = env.step(action_model)
47        
48        # 3. Save to history for charts
49        total_I = sum(env.I)
50        delta_I = total_I - history[-1]["total_infections"] if len(history) > 0 else total_I
51        
52        history.append({
53            "day": env.current_day,
54            "total_infections": total_I,
55            "delta_infections": max(0, delta_I),
56            "poor_economy": env.economy_hit[2],
57            "public_trust": env.public_trust,
58            "reasoning": action_model.reasoning,
59            "policy": action_model.policy_choice
60        })
61
62st.sidebar.divider()
63
64# God Mode Features
65st.sidebar.markdown("### ⚡ God Mode: Reality Injection")
66st.sidebar.caption("Test the AI with real-world shocks.")
67
68# 1-Click Demo Buttons
69col1, col2 = st.sidebar.columns(2)
70if col1.button("🦠 Slum Variant"): 
71    st.session_state.scenario_text = "A deadlier, highly contagious variant mutates in the poor tier."
72if col2.button("📉 Market Crash"): 
73    st.session_state.scenario_text = "A sudden global market crash wipes out the poor tier's remaining wealth."
74
75user_anomaly = st.sidebar.text_input("Describe an event:", value=st.session_state.scenario_text)
76
77if st.sidebar.button("Inject Event", use_container_width=True):
78    if user_anomaly:
79        with st.spinner("AI Translating NLP to Math..."):
80            effect = agent.interpret_anomaly(user_anomaly)
81            env.apply_dynamic_anomaly(effect['target'], effect['multiplier'])
82            st.toast(f"System Shocked! {effect['target'].upper()} altered by {effect['multiplier']}x", icon="🚨")
83
84# --- 4. MAIN DASHBOARD (METRICS) ---
85st.title("🏛️ EpidemicAI Command Center")
86
87m1, m2, m3, m4 = st.columns(4)
88current_trust = env.public_trust
89
90m1.metric("Current Day", f"{env.current_day} / {env.max_days}")
91m2.metric("Total Active Infections", f"{int(sum(env.I)):,}")
92m3.metric("Poor Tier Wealth (Damage)", f"${int(env.economy_hit[2]):,}")
93m4.metric("Public Trust", f"{current_trust:.1f}%", 
94          delta="RIOT WARNING" if current_trust <= 20 else "Stable", 
95          delta_color="inverse" if current_trust <= 20 else "normal")
96
97# --- 5. THE SOCIAL PULSE ---
98if current_trust > 70:
99    st.success("📱 @Citizen123: The Mayor is handling this perfectly. We feel safe! #FlattenTheCurve")
100elif current_trust > 40:
101    st.info("📱 @CityWorker: It's tough, but we are holding on. Hoping the economy opens soon.")
102elif current_trust > 20:
103    st.warning("📱 @LocalOwner: I can't pay rent. If they don't open up, we lose everything! #OpenUp")
104else:
105    st.error("📱 @AngryMob: WE ARE IGNORING THE LOCKDOWN. NO MORE RULES! 🛑🔥 #Riot")
106
107st.divider()
108
109# --- 6. INTERACTIVE DELTA CHARTS ---
110if len(history) > 0:
111    c1, c2 = st.columns(2)
112    
113    # Chart 1: Daily New Infections (Delta)
114    days = [h["day"] for h in history]
115    deltas = [h["delta_infections"] for h in history]
116    fig_inf = px.bar(x=days, y=deltas, labels={"x": "Day", "y": "New Infections"}, title="Daily New Infections (Spikes)")
117    fig_inf.update_traces(marker_color='crimson')
118    c1.plotly_chart(fig_inf, use_container_width=True)
119    
120    # Chart 2: Poor Economy
121    econ = [h["poor_economy"] for h in history]
122    fig_econ = px.line(x=days, y=econ, labels={"x": "Day", "y": "Economic Damage ($)"}, title="Poor Economy Impact vs Day")
123    fig_econ.update_traces(line_color='orange')
124    fig_econ.add_hline(y=3000, line_dash="dash", line_color="red", annotation_text="Bankruptcy")
125    c2.plotly_chart(fig_econ, use_container_width=True)
126
127    # --- 7. THE AI CABINET DEBATE ---
128    st.subheader("🧠 Live AI Cabinet Debate")
129    latest_reasoning = history[-1]["reasoning"]
130    
131    # Make it look beautiful
132    if "|" in latest_reasoning:
133        parts = latest_reasoning.split("|")
134        for part in parts:
135            if "CMO:" in part:
136                st.info(f"**🩺 CMO:** {part.replace('CMO:', '').strip()}")
137            elif "ECON:" in part:
138                st.warning(f"**💼 ECON:** {part.replace('ECON:', '').strip()}")
139            elif "MAYOR:" in part:
140                st.success(f"**⚖️ MAYOR'S DECISION:** {part.replace('MAYOR:', '').strip()}")
141    else:
142        st.write(latest_reasoning)
143else:
144    st.info("👈 Click 'Run Next Day' in the sidebar to begin the simulation.")