Update app.py

app.py

@@ -1,5 +1,5 @@
 # app.py
-# @title Beer Game Final Version (v4.21 - Removed Introduction)
+# @title Beer Game Final Version (v4.30 - Heterogeneous "Locus of Chaos" Design)
 # -----------------------------------------------------------------------------
 # 1. Import Libraries
 # -----------------------------------------------------------------------------
@@ -59,19 +59,44 @@ else:
 def get_customer_demand(week: int) -> int:
     return 4 if week <= 4 else 8
 
-def init_game_state(llm_personality: str, info_sharing: str):
+def init_game_state(locus_of_chaos: str, info_sharing: str):
+    """
+    Initializes the game state based on the Locus of Chaos and Information Sharing conditions.
+    The human role is fixed as 'Distributor'.
+    """
     roles = ["Retailer", "Wholesaler", "Distributor", "Factory"]
-    human_role = "Distributor" # Role is fixed
+    human_role = "Distributor" # Role is fixed as per our discussion
     participant_id = str(uuid.uuid4())[:8]
+
+    # --- NEW: Define heterogeneous personalities based on Locus of Chaos ---
+    if locus_of_chaos == 'Downstream Chaos':
+        personalities = {
+            "Retailer": "human_like",
+            "Wholesaler": "human_like",
+            "Distributor": "HUMAN_PLAYER", # Human role
+            "Factory": "perfect_rational"
+        }
+    else: # 'Upstream Chaos'
+        personalities = {
+            "Retailer": "perfect_rational",
+            "Wholesaler": "perfect_rational",
+            "Distributor": "HUMAN_PLAYER", # Human role
+            "Factory": "human_like"
+        }
+    # ---------------------------------------------------------------------
+
     st.session_state.game_state = {
         'game_running': True, 'participant_id': participant_id, 'week': 1,
-        'human_role': human_role, 'llm_personality': llm_personality,
+        'human_role': human_role,
+        'locus_of_chaos': locus_of_chaos, # <-- NEW: Store chaos condition
         'info_sharing': info_sharing, 'logs': [], 'echelons': {},
         'factory_production_pipeline': deque([0] * FACTORY_LEAD_TIME, maxlen=FACTORY_LEAD_TIME),
         'decision_step': 'initial_order',
         'human_initial_order': None,
         'last_week_orders': {name: 0 for name in roles}
+        # 'llm_personality' is now REMOVED from the global state
     }
+
     for i, name in enumerate(roles):
         upstream = roles[i + 1] if i + 1 < len(roles) else None
         downstream = roles[i - 1] if i - 1 >= 0 else None
@@ -79,18 +104,21 @@ def init_game_state(llm_personality: str, info_sharing: str):
         elif name == "Factory": shipping_weeks = 0
         else: shipping_weeks = SHIPPING_DELAY
         st.session_state.game_state['echelons'][name] = {
-            'name': name, 'inventory': INITIAL_INVENTORY, 'backlog': INITIAL_BACKLOG,
+            'name': name,
+            'personality': personalities[name], # <-- NEW: Store agent-specific personality
+            'inventory': INITIAL_INVENTORY, 'backlog': INITIAL_BACKLOG,
             'incoming_shipments': deque([0] * shipping_weeks, maxlen=shipping_weeks),
             'incoming_order': 0, 'order_placed': 0, 'shipment_sent': 0,
             'weekly_cost': 0, 'total_cost': 0, 'upstream_name': upstream, 'downstream_name': downstream,
         }
-    st.info(f"New game started! AI Mode: **{llm_personality} / {info_sharing}**. You are playing as the: **{human_role}**.")
+    st.info(f"New game started! AI Mode: **{locus_of_chaos} / {info_sharing}**. You are playing as the: **{human_role}**.")
 
 def get_llm_order_decision(prompt: str, echelon_name: str) -> (int, str):
     # This function remains correct.
     if not client: return 8, "NO_API_KEY_DEFAULT"
     with st.spinner(f"Getting AI decision for {echelon_name}..."):
         try:
+            # Use lower temp for rational, higher for human-like
             temp = 0.1 if 'rational' in prompt else 0.7
             response = client.chat.completions.create(
                 model=OPENAI_MODEL,
@@ -110,7 +138,11 @@ def get_llm_order_decision(prompt: str, echelon_name: str) -> (int, str):
             return 4, f"API_ERROR: {e}"
 
 def get_llm_prompt(echelon_state_decision_point: dict, week: int, llm_personality: str, info_sharing: str, all_echelons_state_decision_point: dict) -> str:
-    # This function's logic remains correct (from v4.17).
+    """
+    Generates the prompt for a specific AI agent based on its *individual* personality.
+    NO CHANGE WAS NEEDED in this function's logic, as it correctly routes
+    based on the llm_personality string it receives.
+    """
     e_state = echelon_state_decision_point
     base_info = f"Your Current Status at the **{e_state['name']}** for **Week {week}** (Before Shipping):\n- On-hand inventory: {e_state['inventory']} units.\n- Backlog (total unfilled orders): {e_state['backlog']} units.\n- Incoming order this week (just received): {e_state['incoming_order']} units.\n"
     if e_state['name'] == 'Factory':
@@ -136,7 +168,7 @@ def get_llm_prompt(echelon_state_decision_point: dict, week: int, llm_personalit
             inv_pos_components = f"(Inv={e_state['inventory']} - Backlog={e_state['backlog']} + InTransitShip={sum(e_state['incoming_shipments'])} + OrderToSupplier={order_in_transit_to_supplier})"
         optimal_order = max(0, int(target_inventory_level - inventory_position))
         return f"**You are a perfectly rational supply chain AI with full system visibility.**\nYour only goal is to maintain stability and minimize costs based on mathematical optimization.\n**System Analysis:**\n* **Known Stable End-Customer Demand:** {stable_demand} units/week.\n* **Your Current Total Inventory Position:** {inventory_position} units. {inv_pos_components}\n* **Optimal Target Inventory Level:** {target_inventory_level} units (Target for {total_lead_time} weeks lead time).\n* **Mathematically Optimal {task_word.title()}:** The optimal decision is **{optimal_order} units**.\n**Your Task:** Confirm this optimal {task_word}. Respond with a single integer."
-
+    
     elif llm_personality == 'perfect_rational' and info_sharing == 'local':
         safety_stock = 4; anchor_demand = e_state['incoming_order']
         inventory_correction = safety_stock - (e_state['inventory'] - e_state['backlog'])
@@ -150,7 +182,7 @@ def get_llm_prompt(echelon_state_decision_point: dict, week: int, llm_personalit
         calculated_order = anchor_demand + inventory_correction - supply_line
         rational_local_order = max(0, int(calculated_order))
         return f"**You are a perfectly rational supply chain AI with ONLY LOCAL information.**\nYou must use a logical heuristic to make a stable decision. A proven method is \"Anchoring and Adjustment\".\n\n{base_info}\n\n**Rational Calculation (Anchoring & Adjustment):**\n1.  **Anchor on Demand:** Your best guess for future demand is your last incoming order: **{anchor_demand} units**.\n2.  **Adjust for Inventory:** You want to hold a safety stock of {safety_stock} units. Your current stock (before shipping) is {e_state['inventory'] - e_state['backlog']}. You need to order an extra **{inventory_correction} units** to correct this.\n3.  **Account for {supply_line_desc}:** You already have **{supply_line} units** being processed. These should be subtracted from your new decision.\n\n**Final Calculation:**\n* Decision = (Anchor Demand) + (Inventory Adjustment) - ({supply_line_desc})\n* Decision = {anchor_demand} + {inventory_correction} - {supply_line} = **{rational_local_order} units**.\n**Your Task:** Confirm this locally rational {task_word}. Respond with a single integer."
-
+    
     elif llm_personality == 'human_like' and info_sharing == 'full':
         full_info_str = f"\n**Full Supply Chain Information (State Before Shipping):**\n- End-Customer Demand this week: {get_customer_demand(week)} units.\n"
         for name, other_e_state in all_echelons_state_decision_point.items():
@@ -165,7 +197,7 @@ def get_llm_prompt(echelon_state_decision_point: dict, week: int, llm_personalit
         You are still human and might get anxious about your own stock levels.
         What {task_word} should you decide on this week? Respond with a single integer.
         """
-
+    
     elif llm_personality == 'human_like' and info_sharing == 'local':
         return f"""
         **You are a reactive supply chain manager for the {e_state['name']}.** You have a limited view and tend to over-correct based on fear.
@@ -177,33 +209,33 @@ def get_llm_prompt(echelon_state_decision_point: dict, week: int, llm_personalit
         """
 
 def step_game(human_final_order: int, human_initial_order: int, ai_suggestion: int):
-    # This function's logic remains correct (from v4.17).
+    # This function's core logic remains correct.
     state = st.session_state.game_state
     week, echelons, human_role = state['week'], state['echelons'], state['human_role']
-    llm_personality, info_sharing = state['llm_personality'], state['info_sharing']
+    # --- MODIFIED: Get info_sharing, but llm_personality is now per-echelon ---
+    info_sharing = state['info_sharing']
     echelon_order = ["Retailer", "Wholesaler", "Distributor", "Factory"]
+    
     llm_raw_responses = {}
-
     opening_inventories = {name: e['inventory'] for name, e in echelons.items()}
     opening_backlogs = {name: e['backlog'] for name, e in echelons.items()}
-
     arrived_this_week = {name: 0 for name in echelon_order}
     inventory_after_arrival = {}
-
+    
     factory_state = echelons["Factory"]
     produced_units = 0
     if state['factory_production_pipeline']:
         produced_units = state['factory_production_pipeline'].popleft()
         arrived_this_week["Factory"] = produced_units
     inventory_after_arrival["Factory"] = factory_state['inventory'] + produced_units
-
+    
     for name in ["Retailer", "Wholesaler", "Distributor"]:
         arrived_shipment = 0
         if echelons[name]['incoming_shipments']:
             arrived_shipment = echelons[name]['incoming_shipments'].popleft()
             arrived_this_week[name] = arrived_shipment
         inventory_after_arrival[name] = echelons[name]['inventory'] + arrived_shipment
-
+    
     total_backlog_before_shipping = {}
     for name in echelon_order:
         incoming_order_for_this_week = 0
@@ -213,7 +245,7 @@ def step_game(human_final_order: int, human_initial_order: int, ai_suggestion: i
             if downstream_name: incoming_order_for_this_week = state['last_week_orders'].get(downstream_name, 0)
         echelons[name]['incoming_order'] = incoming_order_for_this_week
         total_backlog_before_shipping[name] = echelons[name]['backlog'] + incoming_order_for_this_week
-
+    
     decision_point_states = {}
     for name in echelon_order:
          decision_point_states[name] = {
@@ -221,42 +253,58 @@ def step_game(human_final_order: int, human_initial_order: int, ai_suggestion: i
              'backlog': total_backlog_before_shipping[name], 'incoming_order': echelons[name]['incoming_order'],
              'incoming_shipments': echelons[name]['incoming_shipments'].copy() if name != "Factory" else deque(),
          }
-
+    
     current_week_orders = {}
     for name in echelon_order:
         e = echelons[name]; prompt_state = decision_point_states[name]
-        if name == human_role: order_amount, raw_resp = human_final_order, "HUMAN_FINAL_INPUT"
+        
+        if name == human_role: 
+            order_amount, raw_resp = human_final_order, "HUMAN_FINAL_INPUT"
         else:
-            prompt = get_llm_prompt(prompt_state, week, llm_personality, info_sharing, decision_point_states)
+            # --- MODIFIED: Get the specific agent's personality ---
+            e_personality = e['personality'] 
+            prompt = get_llm_prompt(prompt_state, week, e_personality, info_sharing, decision_point_states)
             order_amount, raw_resp = get_llm_order_decision(prompt, name)
+            
         llm_raw_responses[name] = raw_resp; e['order_placed'] = max(0, order_amount); current_week_orders[name] = e['order_placed']
-
+    
     state['factory_production_pipeline'].append(echelons["Factory"]['order_placed'])
-
     units_shipped = {name: 0 for name in echelon_order}
+    
     for name in echelon_order:
         e = echelons[name]; demand_to_meet = total_backlog_before_shipping[name]; available_inv = inventory_after_arrival[name]
         e['shipment_sent'] = min(available_inv, demand_to_meet); units_shipped[name] = e['shipment_sent']
         e['inventory'] = available_inv - e['shipment_sent']; e['backlog'] = demand_to_meet - e['shipment_sent']
-
+    
     if units_shipped["Factory"] > 0: echelons['Distributor']['incoming_shipments'].append(units_shipped["Factory"])
     if units_shipped['Distributor'] > 0: echelons['Wholesaler']['incoming_shipments'].append(units_shipped['Distributor'])
     if units_shipped['Wholesaler'] > 0: echelons['Retailer']['incoming_shipments'].append(units_shipped['Wholesaler'])
-
+    
+    # --- MODIFIED: Update logging fields ---
     log_entry = {'timestamp': datetime.utcnow().isoformat() + "Z", 'week': week, **state}
+    # Remove large state objects from log entry
     del log_entry['echelons'], log_entry['factory_production_pipeline'], log_entry['logs'], log_entry['last_week_orders']
+    # 'llm_personality' is already gone from state
+    
     for name in echelon_order:
         e = echelons[name]; e['weekly_cost'] = (e['inventory'] * HOLDING_COST) + (e['backlog'] * BACKLOG_COST); e['total_cost'] += e['weekly_cost']
-        for key in ['inventory', 'backlog', 'incoming_order', 'order_placed', 'shipment_sent', 'weekly_cost', 'total_cost']: log_entry[f'{name}.{key}'] = e[key]
+        for key in ['inventory', 'backlog', 'incoming_order', 'order_placed', 'shipment_sent', 'weekly_cost', 'total_cost']: 
+            log_entry[f'{name}.{key}'] = e[key]
+        
+        log_entry[f'{name}.personality'] = e['personality'] # <-- NEW: Log individual personality
         log_entry[f'{name}.llm_raw_response'] = llm_raw_responses.get(name, "")
         log_entry[f'{name}.opening_inventory'] = opening_inventories[name]; log_entry[f'{name}.opening_backlog'] = opening_backlogs[name]
         log_entry[f'{name}.arrived_this_week'] = arrived_this_week[name]
-        if name != 'Factory': log_entry[f'{name}.arriving_next_week'] = list(e['incoming_shipments'])[0] if e['incoming_shipments'] else 0
-        else: log_entry[f'{name}.production_completing_next_week'] = list(state['factory_production_pipeline'])[0] if state['factory_production_pipeline'] else 0
+        if name != 'Factory': 
+            log_entry[f'{name}.arriving_next_week'] = list(e['incoming_shipments'])[0] if e['incoming_shipments'] else 0
+        else: 
+            log_entry[f'{name}.production_completing_next_week'] = list(state['factory_production_pipeline'])[0] if state['factory_production_pipeline'] else 0
+            
     log_entry[f'{human_role}.initial_order'] = human_initial_order; log_entry[f'{human_role}.ai_suggestion'] = ai_suggestion
+    
     state['logs'].append(log_entry)
-
     state['week'] += 1; state['decision_step'] = 'initial_order'; state['last_week_orders'] = current_week_orders
+    
     if state['week'] > WEEKS: state['game_running'] = False
 
 def plot_results(df: pd.DataFrame, title: str, human_role: str):
@@ -264,6 +312,7 @@ def plot_results(df: pd.DataFrame, title: str, human_role: str):
     fig, axes = plt.subplots(4, 1, figsize=(12, 22))
     fig.suptitle(title, fontsize=16)
     echelons = ['Retailer', 'Wholesaler', 'Distributor', 'Factory']
+    
     plot_data = []
     for _, row in df.iterrows():
         for e in echelons:
@@ -271,13 +320,17 @@ def plot_results(df: pd.DataFrame, title: str, human_role: str):
                 'inventory': row.get(f'{e}.inventory', 0), 'order_placed': row.get(f'{e}.order_placed', 0),
                 'total_cost': row.get(f'{e}.total_cost', 0)})
     plot_df = pd.DataFrame(plot_data)
+    
     inventory_pivot = plot_df.pivot(index='week', columns='echelon', values='inventory').reindex(columns=echelons)
     inventory_pivot.plot(ax=axes[0], kind='line', marker='o', markersize=4); axes[0].set_title('Inventory Levels (End of Week)'); axes[0].grid(True, linestyle='--'); axes[0].set_ylabel('Stock (Units)')
+    
     order_pivot = plot_df.pivot(index='week', columns='echelon', values='order_placed').reindex(columns=echelons)
     order_pivot.plot(ax=axes[1], style='--'); axes[1].plot(range(1, WEEKS + 1), [get_customer_demand(w) for w in range(1, WEEKS + 1)], label='Customer Demand', color='black', lw=2.5); axes[1].set_title('Order Quantities / Production Decisions'); axes[1].grid(True, linestyle='--'); axes[1].legend(); axes[1].set_ylabel('Ordered/Produced (Units)')
+    
     total_costs = plot_df.loc[plot_df.groupby('echelon')['week'].idxmax()]
     total_costs = total_costs.set_index('echelon')['total_cost'].reindex(echelons, fill_value=0)
     total_costs.plot(kind='bar', ax=axes[2], rot=0); axes[2].set_title('Total Cumulative Cost'); axes[2].set_ylabel('Cost ($)')
+    
     human_cols = [f'{human_role}.initial_order', f'{human_role}.ai_suggestion', f'{human_role}.order_placed']
     human_df_cols = ['week'] + [col for col in human_cols if col in df.columns]
     try:
@@ -287,6 +340,7 @@ def plot_results(df: pd.DataFrame, title: str, human_role: str):
         else: raise ValueError("No human decision data columns found.")
     except (KeyError, ValueError) as plot_err:
         axes[3].set_title(f'Analysis of Your ({human_role}) Decisions - Error Plotting Data'); axes[3].text(0.5, 0.5, f"Error: {plot_err}", ha='center', va='center'); axes[3].grid(True, linestyle='--'); axes[3].set_xlabel('Week')
+    
     plt.tight_layout(rect=[0, 0, 1, 0.96]); return fig
 
 def save_logs_and_upload(state: dict):
@@ -319,56 +373,64 @@ else:
     # --- Game Setup & Instructions ---
     if 'game_state' not in st.session_state or not st.session_state.game_state.get('game_running', False):
         
-        # --- Introduction Section Removed as Requested ---
-        
         st.header("⚙️ Game Configuration")
         c1, c2 = st.columns(2)
         with c1:
-             llm_personality = st.selectbox("AI Agent 'Personality'", ('human_like', 'perfect_rational'), format_func=lambda x: x.replace('_', ' ').title(), help="**Human-like:** Tends to react emotionally, potentially over-ordering. **Perfect Rational:** Uses a mathematical heuristic to make stable, logical decisions.")
+            # --- MODIFIED: Changed from llm_personality to locus_of_chaos ---
+            locus_of_chaos = st.selectbox(
+                "AI Team Composition (Locus of Chaos)", 
+                ('Downstream Chaos', 'Upstream Chaos'), 
+                format_func=lambda x: x.replace('_', ' ').title(), 
+                help=(
+                    "**Downstream Chaos:** Your customers (Retailer, Wholesaler) are 'Human-like' (chaotic). Your supplier (Factory) is 'Rational'.\n\n"
+                    "**Upstream Chaos:** Your customers (Retailer, Wholesaler) are 'Rational' (stable). Your supplier (Factory) is 'Human-like'."
+                )
+            )
         with c2:
-            info_sharing = st.selectbox("Information Sharing Level", ('local', 'full'), format_func=lambda x: x.title(), help="**Local:** You and the AI agents can only see your own inventory and incoming orders. **Full:** Everyone can see the entire supply chain's status and the true end-customer demand.")
-        
+            info_sharing = st.selectbox(
+                "Information Sharing Level", 
+                ('local', 'full'), 
+                format_func=lambda x: x.title(), 
+                help="**Local:** You and the AI agents can only see your own inventory and incoming orders. **Full:** Everyone can see the entire supply chain's status and the true end-customer demand."
+            )
+                
         if st.button("🚀 Start Game", type="primary", disabled=(client is None)):
-            init_game_state(llm_personality, info_sharing)
+            # --- MODIFIED: Pass locus_of_chaos to init ---
+            init_game_state(locus_of_chaos, info_sharing)
             st.rerun()
 
     # --- Main Game Interface ---
     elif 'game_state' in st.session_state and st.session_state.game_state.get('game_running'):
         state = st.session_state.game_state
         week, human_role, echelons, info_sharing = state['week'], state['human_role'], state['echelons'], state['info_sharing']
-        echelon_order = ["Retailer", "Wholesaler", "Distributor", "Factory"] # Define here for UI
-        
+        echelon_order = ["Retailer", "Wholesaler", "Distributor", "Factory"]
+                
         st.header(f"Week {week} / {WEEKS}")
-        st.subheader(f"Your Role: **{human_role}** | AI Mode: **{state['llm_personality'].replace('_', ' ')}** | Information: **{state['info_sharing']}**")
+        # --- MODIFIED: Update subheader to use locus_of_chaos ---
+        st.subheader(f"Your Role: **{human_role}** | AI Mode: **{state['locus_of_chaos']}** | Information: **{state['info_sharing']}**")
         st.markdown("---")
-        
-        st.subheader("Supply Chain Status (Start of Week State)") # Clarified Timing
-        
+                
+        st.subheader("Supply Chain Status (Start of Week State)")
+                
         if info_sharing == 'full':
             cols = st.columns(4)
-            for i, name in enumerate(echelon_order): # Use the defined echelon_order
+            for i, name in enumerate(echelon_order):
                 with cols[i]:
-                    e = echelons[name] # Get the echelon state
+                    e = echelons[name] 
                     icon = "👤" if name == human_role else "🤖"
-                    
-                    # =============== UI CHANGE: Highlight Player ===============
+                                        
                     if name == human_role:
-                        # Use markdown with HTML/CSS for highlighting
                         st.markdown(f"##### **<span style='border: 1px solid #FF4B4B; padding: 2px 5px; border-radius: 3px;'>{icon} {name} (You)</span>**", unsafe_allow_html=True)
                     else:
                         st.markdown(f"##### {icon} {name}")
-                    # ========================================================
-                    
-                    # Display the END OF LAST WEEK state (which is OPENING state for this week)
+                        # --- NEW: Display the specific AI's personality ---
+                        personality_label = e['personality'].replace('_', ' ').title()
+                        st.caption(f"AI Type: **{personality_label}**")
+                        # -------------------------------------------------
+                                        
                     st.metric("Inventory (Opening)", e['inventory'])
                     st.metric("Backlog (Opening)", e['backlog'])
-                    
-                    # =============== UI CHANGE: Removed Costs ===============
-                    # Costs are no longer displayed on the main dashboard
-                    # =======================================================
-                    
-                    # Display info about THIS week's events / NEXT week's arrivals
-                    # Calculate the INCOMING order for THIS week
+                                        
                     current_incoming_order = 0
                     if name == "Retailer":
                         current_incoming_order = get_customer_demand(week)
@@ -376,111 +438,121 @@ else:
                         downstream_name = e['downstream_name']
                         if downstream_name:
                             current_incoming_order = state['last_week_orders'].get(downstream_name, 0)
-                    
-                    st.write(f"Incoming Order (This Week): **{current_incoming_order}**") # Display calculated order
-                    
-                    # Display prediction for NEXT week's arrivals
+                                        
+                    st.write(f"Incoming Order (This Week): **{current_incoming_order}**") 
+                                        
                     if name == "Factory":
                         prod_completing_next = list(state['factory_production_pipeline'])[0] if state['factory_production_pipeline'] else 0
                         st.write(f"Completing Next Week: **{prod_completing_next}**")
                     else:
                         arriving_next = list(e['incoming_shipments'])[0] if e['incoming_shipments'] else 0
                         st.write(f"Arriving Next Week: **{arriving_next}**")
+        
         else: # Local Info Mode
             st.info("In Local Information mode, you can only see your own status dashboard.")
             e = echelons[human_role]
-            st.markdown(f"### 👤 **<span style='color:#FF4B4B;'>{human_role} (Your Dashboard - Start of Week State)</span>**", unsafe_allow_html=True) # Highlight self
+            st.markdown(f"### 👤 **<span style='color:#FF4B4B;'>{human_role} (Your Dashboard - Start of Week State)</span>**", unsafe_allow_html=True)
             col1, col2, col3, col4 = st.columns(4)
             
-            # Display OPENING state
             col1.metric("Inventory (Opening)", e['inventory'])
             col2.metric("Backlog (Opening)", e['backlog'])
-            
-            # Display info about THIS week's events / NEXT week's arrivals
-            # Calculate the INCOMING order for THIS week
+                        
             current_incoming_order = 0
             downstream_name = e['downstream_name'] # Wholesaler
-            if downstream_name:
+            if downstream_name: 
                  current_incoming_order = state['last_week_orders'].get(downstream_name, 0)
-            
-            col3.write(f"**Incoming Order (This Week):**\n# {current_incoming_order}") # Display calculated order
+                        
+            col3.write(f"**Incoming Order (This Week):**\n# {current_incoming_order}")
             col4.write(f"**Shipment Arriving (Next Week):**\n# {list(e['incoming_shipments'])[0] if e['incoming_shipments'] else 0}")
             
-            # =============== UI CHANGE: Removed Costs ===============
-            # Costs are no longer displayed on the main dashboard
-            # =======================================================
-
         st.markdown("---")
         st.header("Your Decision (Step 4)")
-
+        
         # Prepare the state snapshot for the AI prompt (State AFTER arrivals/orders, BEFORE shipping)
         all_decision_point_states = {}
         for name in echelon_order:
-            e_curr = echelons[name] # This is END OF LAST WEEK state
+            e_curr = echelons[name] 
             arrived = 0
-            # Peek at what *will* arrive this week (Step 1) based on current queues
-            if name == "Factory":
+            if name == "Factory": 
                  if state['factory_production_pipeline']: arrived = list(state['factory_production_pipeline'])[0]
-            else:
+            else: 
                  if e_curr['incoming_shipments']: arrived = list(e_curr['incoming_shipments'])[0]
-            
-            # Calculate the state AFTER arrivals and incoming orders for the prompt
+                        
             inv_after_arrival = e_curr['inventory'] + arrived
-            
-            # Determine incoming order for *this* week again for prompt state
+                        
             inc_order_this_week = 0
             if name == "Retailer": inc_order_this_week = get_customer_demand(week)
             else:
                 ds_name = e_curr['downstream_name']
                 if ds_name: inc_order_this_week = state['last_week_orders'].get(ds_name, 0)
-            
+                        
             backlog_after_new_order = e_curr['backlog'] + inc_order_this_week
-            
+                        
             all_decision_point_states[name] = {
                 'name': name, 'inventory': inv_after_arrival, 'backlog': backlog_after_new_order,
-                'incoming_order': inc_order_this_week, # Use the correctly calculated incoming order
+                'incoming_order': inc_order_this_week, 
                 'incoming_shipments': e_curr['incoming_shipments'].copy() if name != "Factory" else deque()
             }
-        
+                
         human_echelon_state_for_prompt = all_decision_point_states[human_role]
-
+        
         if state['decision_step'] == 'initial_order':
             with st.form(key="initial_order_form"):
                 st.markdown("#### **Step 4a:** Based on the dashboard, submit your **initial** order to the Factory.")
-                # =============== UI CHANGE: Removed Default Value ===============
-                initial_order = st.number_input("Your Initial Order Quantity:", min_value=0, step=1) # No 'value' argument
-                # ===============================================================
+                initial_order = st.number_input("Your Initial Order Quantity:", min_value=0, step=1)
                 if st.form_submit_button("Submit Initial Order & See AI Suggestion", type="primary"):
-                    # Handle case where user leaves it blank (input returns None)
                     state['human_initial_order'] = int(initial_order) if initial_order is not None else 0
                     state['decision_step'] = 'final_order'
                     st.rerun()
-
+        
         elif state['decision_step'] == 'final_order':
             st.success(f"Your initial order was: **{state['human_initial_order']}** units.")
+                        
+            # --- MODIFIED: Get the human's "partner" AI personality for the suggestion ---
+            # In our design, the human (Distributor) gets a suggestion from an AI *acting as* the Distributor.
+            # We must decide which personality this "suggestion AI" should have.
+            # For simplicity, we'll use the personality defined for the HUMAN'S ROLE in the `personalities` dict.
+            # ...wait, that's "HUMAN_PLAYER".
+            #
+            # --- CORRECTION / EXECUTIVE DECISION ---
+            # The *suggestion* AI should match the human's role. But what personality?
+            # Let's assume the "Suggestion AI" is a *separate* entity that matches the *dominant* mode of the other AIs.
+            # This is complex.
+            #
+            # --- SIMPLER, BETTER LOGIC ---
+            # The experiment is about interacting with AI. The human *is* the Distributor.
+            # The AI *suggestion* should come from an AI also *simulating* the Distributor role.
+            # What personality should it have? 
+            # Let's make the suggestion AI's personality *also* dependent on the Locus of Chaos.
+            # In 'Downstream Chaos', the human is surrounded by 'human_like' AIs. Their suggestion should be 'human_like'.
+            # In 'Upstream Chaos', the human is surrounded by 'perfect_rational' AIs. Their suggestion should be 'perfect_rational'.
+            #
+            # The human (Distributor)'s customers are Retailer/Wholesaler.
+            # So, the "suggestion" AI's personality will match the personality of the human's *customers*.
             
-            # Use the correctly timed state for the prompt
-            prompt_sugg = get_llm_prompt(human_echelon_state_for_prompt, week, state['llm_personality'], state['info_sharing'], all_decision_point_states)
-            ai_suggestion, _ = get_llm_order_decision(prompt_sugg, f"{human_role} (Suggestion)")
+            if state['locus_of_chaos'] == 'Downstream Chaos':
+                suggestion_ai_personality = 'human_like' # Matches chaotic customers
+            else: # 'Upstream Chaos'
+                suggestion_ai_personality = 'perfect_rational' # Matches rational customers
+            # ------------------------------------------------
             
+            prompt_sugg = get_llm_prompt(human_echelon_state_for_prompt, week, suggestion_ai_personality, state['info_sharing'], all_decision_point_states)
+            ai_suggestion, _ = get_llm_order_decision(prompt_sugg, f"{human_role} (Suggestion)")
+                        
             with st.form(key="final_order_form"):
-                st.markdown(f"#### **Step 4b:** The AI suggests ordering **{ai_suggestion}** units.")
+                st.markdown(f"#### **Step 4b:** An AI {suggestion_ai_personality.replace('_', ' ')} assistant suggests ordering **{ai_suggestion}** units.")
                 st.markdown("Considering the AI's advice, submit your **final** order to end the week. (This order will arrive in 3 weeks).")
-                # =============== UI CHANGE: Removed Default Value ===============
-                st.number_input("Your Final Order Quantity:", min_value=0, step=1, key='final_order_input') # No 'value' argument
-                # ===============================================================
-                
+                st.number_input("Your Final Order Quantity:", min_value=0, step=1, key='final_order_input') 
+                                
                 if st.form_submit_button("Submit Final Order & Advance to Next Week"):
-                    # Handle case where user leaves it blank
-                    final_order_value = st.session_state.get('final_order_input', 0) # Use .get with default
+                    final_order_value = st.session_state.get('final_order_input', 0) 
                     final_order_value = int(final_order_value) if final_order_value is not None else 0
-                    
+                                        
                     step_game(final_order_value, state['human_initial_order'], ai_suggestion)
-                    
-                    # Clean up session state for the input key
+                                        
                     if 'final_order_input' in st.session_state: del st.session_state.final_order_input
                     st.rerun()
-
+                    
         st.markdown("---")
         with st.expander("📖 Your Weekly Decision Log", expanded=False):
             if not state.get('logs'):
@@ -499,10 +571,10 @@ else:
                         'week', f'{human_role}.opening_inventory', f'{human_role}.opening_backlog',
                         f'{human_role}.arrived_this_week', f'{human_role}.incoming_order',
                         f'{human_role}.initial_order', f'{human_role}.ai_suggestion', f'{human_role}.order_placed',
-                        f'{human_role}.arriving_next_week', f'{human_role}.weekly_cost'
+                        f'{human_role}.arriving_next_week', f'{human_role}.weekly_cost' 
                      ]
                     final_cols_to_display = [col for col in ordered_display_cols_keys if col in history_df.columns]
-                    if not final_cols_to_display:
+                    if not final_cols_to_display: 
                          st.write("No data columns available to display.")
                     else:
                         display_df = history_df[final_cols_to_display].rename(columns=human_cols)
@@ -511,7 +583,7 @@ else:
                         st.dataframe(display_df.sort_values(by="Week", ascending=False), hide_index=True, use_container_width=True)
                 except Exception as e:
                     st.error(f"Error displaying weekly log: {e}")
-
+                    
         try: st.sidebar.image(IMAGE_PATH, caption="Supply Chain Reference")
         except FileNotFoundError: st.sidebar.warning("Image file not found.")
         st.sidebar.header("Game Info")
@@ -527,16 +599,17 @@ else:
         state = st.session_state.game_state
         try:
             logs_df = pd.json_normalize(state['logs'])
+            # --- MODIFIED: Update plot title ---
             fig = plot_results(
                 logs_df,
-                f"Beer Game (Human: {state['human_role']})\n(AI: {state['llm_personality']} | Info: {state['info_sharing']})",
+                f"Beer Game (Human: {state['human_role']})\n(AI Mode: {state['locus_of_chaos']} | Info: {state['info_sharing']})",
                 state['human_role']
             )
             st.pyplot(fig)
             save_logs_and_upload(state)
         except Exception as e:
             st.error(f"Error generating final report: {e}")
-
+            
         if st.button("✨ Start a New Game"):
             del st.session_state.game_state
             st.rerun()