app.py

# app.py
# @title Beer Game Final Version (v4.30 - Heterogeneous "Locus of Chaos" Design)
# -----------------------------------------------------------------------------
# 1. Import Libraries
# -----------------------------------------------------------------------------
import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from collections import deque
import time
import openai
import re
import random
import uuid
from pathlib import Path
from datetime import datetime
from huggingface_hub import HfApi

# -----------------------------------------------------------------------------
# 0. Page Configuration (Must be the first Streamlit command)
# -----------------------------------------------------------------------------
st.set_page_config(page_title="Beer Game: Human-AI Collaboration", layout="wide")

# -----------------------------------------------------------------------------
# 2. Game Parameters & API Configuration
# -----------------------------------------------------------------------------
# --- Game Parameters ---
WEEKS = 24
INITIAL_INVENTORY = 12
INITIAL_BACKLOG = 0
ORDER_PASSING_DELAY = 1 # Handled by last_week_orders
SHIPPING_DELAY = 2 # General shipping delay (R->W, W->D)
FACTORY_LEAD_TIME = 1
FACTORY_SHIPPING_DELAY = 1 # Specific delay from Factory to Distributor
HOLDING_COST = 0.5
BACKLOG_COST = 1.0

# --- Model & Log Configuration ---
OPENAI_MODEL = "gpt-4o-mini"
LOCAL_LOG_DIR = Path("logs")
LOCAL_LOG_DIR.mkdir(exist_ok=True)
IMAGE_PATH = "beer_game_diagram.png" # Path to your uploaded image

# --- API & Secrets Configuration ---
try:
    client = openai.OpenAI(api_key=st.secrets["OPENAI_API_KEY"])
    HF_TOKEN = st.secrets.get("HF_TOKEN")
    HF_REPO_ID = st.secrets.get("HF_REPO_ID")
    hf_api = HfApi() if HF_TOKEN else None
except Exception as e:
    st.session_state.initialization_error = f"Error reading secrets on startup: {e}."
    client = None
else:
    st.session_state.initialization_error = None

# -----------------------------------------------------------------------------
# 3. Core Game Logic Functions
# -----------------------------------------------------------------------------
def get_customer_demand(week: int) -> int:
    return 4 if week <= 4 else 8

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 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,
        '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
        if name == "Distributor": shipping_weeks = FACTORY_SHIPPING_DELAY
        elif name == "Factory": shipping_weeks = 0
        else: shipping_weeks = SHIPPING_DELAY
        st.session_state.game_state['echelons'][name] = {
            '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: **{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,
                messages=[
                    {"role": "system", "content": "You are a supply chain manager playing the Beer Game. Your response must be only an integer number representing your order or production quantity and nothing else. For example: 8"},
                    {"role": "user", "content": prompt}
                ],
                temperature=temp, max_tokens=10
            )
            raw_text = response.choices[0].message.content.strip()
            match = re.search(r'\d+', raw_text)
            if match: return int(match.group(0)), raw_text
            st.warning(f"LLM for {echelon_name} did not return a valid number. Defaulting to 4. Raw Response: '{raw_text}'")
            return 4, raw_text # Default to 4
        except Exception as e:
            st.error(f"API call failed for {echelon_name}: {e}. Defaulting to 4.")
            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:
    """
    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':
        task_word = "production quantity"
        base_info += f"- Your Production Pipeline (completing next week onwards): {list(st.session_state.game_state['factory_production_pipeline'])}"
    else:
        task_word = "order quantity"
        base_info += f"- Shipments In Transit To You (arriving next week onwards): {list(e_state['incoming_shipments'])}"

    if llm_personality == 'perfect_rational' and info_sharing == 'full':
        stable_demand = 8
        if e_state['name'] == 'Factory': total_lead_time = FACTORY_LEAD_TIME
        elif e_state['name'] == 'Distributor': total_lead_time = ORDER_PASSING_DELAY + FACTORY_LEAD_TIME + FACTORY_SHIPPING_DELAY
        else: total_lead_time = ORDER_PASSING_DELAY + SHIPPING_DELAY
        safety_stock = 4
        target_inventory_level = (stable_demand * total_lead_time) + safety_stock
        if e_state['name'] == 'Factory':
            inventory_position = (e_state['inventory'] - e_state['backlog'] + sum(st.session_state.game_state['factory_production_pipeline']))
            inv_pos_components = f"(Inv={e_state['inventory']} - Backlog={e_state['backlog']} + InProd={sum(st.session_state.game_state['factory_production_pipeline'])})"
        else:
            order_in_transit_to_supplier = st.session_state.game_state['last_week_orders'].get(e_state['name'], 0)
            inventory_position = (e_state['inventory'] - e_state['backlog'] + sum(e_state['incoming_shipments']) + order_in_transit_to_supplier)
            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'])
        if e_state['name'] == 'Factory':
            supply_line = sum(st.session_state.game_state['factory_production_pipeline'])
            supply_line_desc = "In Production"
        else:
            order_in_transit_to_supplier = st.session_state.game_state['last_week_orders'].get(e_state['name'], 0)
            supply_line = sum(e_state['incoming_shipments']) + order_in_transit_to_supplier
            supply_line_desc = "Supply Line (In Transit Shipments + Order To Supplier)"
        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():
            if name != e_state['name']: full_info_str += f"- {name}: Inv={other_e_state['inventory']}, Backlog={other_e_state['backlog']}\n"
        return f"""
        **You are a supply chain manager ({e_state['name']}) with full system visibility.**
        You can see everyone's current inventory and backlog before shipping, and the real customer demand.
        {base_info}
        {full_info_str}
        **Your Task:** Your primary responsibility is to meet the demand from your direct customer (your `Incoming order this week`: **{e_state['incoming_order']}** units), which contributes to your total current backlog of {e_state['backlog']}.
        While you can see the stable end-customer demand ({get_customer_demand(week)} units), your priority is to fulfill the order you just received and manage your inventory/backlog.
        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.
        Your top priority is to NOT have a backlog.
        {base_info}
        **Your Task:** You just received an incoming order for **{e_state['incoming_order']}** units, adding to your total backlog.
        Your gut instinct is to panic and {task_word.split(' ')[0]} enough to ensure you are never caught with a backlog again, considering your current inventory.
        **React emotionally.** What is your knee-jerk {task_word}? Respond with a single integer.
        """

def step_game(human_final_order: int, human_initial_order: int, ai_suggestion: int):
    # This function's core logic remains correct.
    state = st.session_state.game_state
    week, echelons, human_role = state['week'], state['echelons'], state['human_role']
    # --- 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
        if name == "Retailer": incoming_order_for_this_week = get_customer_demand(week)
        else:
            downstream_name = echelons[name]['downstream_name']
            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] = {
             'name': name, 'inventory': inventory_after_arrival[name],
             '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"
        else:
            # --- 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]
        
        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
            
    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):
    # This function remains correct.
    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:
            plot_data.append({'week': row.get('week', 0), 'echelon': e,
                '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:
        human_df = df[human_df_cols].copy()
        human_df.rename(columns={ f'{human_role}.initial_order': 'Your Initial Order', f'{human_role}.ai_suggestion': 'AI Suggestion', f'{human_role}.order_placed': 'Your Final Order'}, inplace=True)
        if len(human_df.columns) > 1: human_df.plot(x='week', ax=axes[3], marker='o', linestyle='-'); axes[3].set_title(f'Analysis of Your ({human_role}) Decisions'); axes[3].set_ylabel('Order Quantity'); axes[3].grid(True, linestyle='--'); axes[3].set_xlabel('Week')
        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):
    # This function remains correct.
    if not state.get('logs'): return
    participant_id = state['participant_id']
    try:
        df = pd.json_normalize(state['logs'])
        fname = LOCAL_LOG_DIR / f"log_{participant_id}_{int(time.time())}.csv"
        for col in df.select_dtypes(include=['object']).columns: df[col] = df[col].astype(str)
        df.to_csv(fname, index=False)
        st.success(f"Log successfully saved locally: `{fname}`")
        with open(fname, "rb") as f: st.download_button("📥 Download Log CSV", data=f, file_name=fname.name, mime="text/csv")
        if HF_TOKEN and HF_REPO_ID and hf_api:
            with st.spinner("Uploading log to Hugging Face Hub..."):
                try:
                    url = hf_api.upload_file( path_or_fileobj=str(fname), path_in_repo=f"logs/{fname.name}", repo_id=HF_REPO_ID, repo_type="dataset", token=HF_TOKEN)
                    st.success(f"✅ Log successfully uploaded to Hugging Face! [View File]({url})")
                except Exception as e_upload: st.error(f"Upload to Hugging Face failed: {e_upload}")
    except Exception as e_save: st.error(f"Error processing or saving log data: {e_save}")

# -----------------------------------------------------------------------------
# 4. Streamlit UI (Adjusted Dashboard Labels & Logic)
# -----------------------------------------------------------------------------
st.title("🍺 The Beer Game: A Human-AI Collaboration Challenge")

if st.session_state.get('initialization_error'):
    st.error(st.session_state.initialization_error)
else:
    # --- Game Setup & Instructions ---
    if 'game_state' not in st.session_state or not st.session_state.game_state.get('game_running', False):
        
        st.header("⚙️ Game Configuration")
        c1, c2 = st.columns(2)
        with c1:
            # --- 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."
            )
                
        if st.button("🚀 Start Game", type="primary", disabled=(client is None)):
            # --- 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"]
                
        st.header(f"Week {week} / {WEEKS}")
        # --- 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)")
                
        if info_sharing == 'full':
            cols = st.columns(4)
            for i, name in enumerate(echelon_order):
                with cols[i]:
                    e = echelons[name] 
                    icon = "👤" if name == human_role else "🤖"
                                        
                    if name == human_role:
                        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}")
                        # --- 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'])
                                        
                    current_incoming_order = 0
                    if name == "Retailer":
                        current_incoming_order = get_customer_demand(week)
                    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}**") 
                                        
                    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)
            col1, col2, col3, col4 = st.columns(4)
            
            col1.metric("Inventory (Opening)", e['inventory'])
            col2.metric("Backlog (Opening)", e['backlog'])
                        
            current_incoming_order = 0
            downstream_name = e['downstream_name'] # Wholesaler
            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}")
            col4.write(f"**Shipment Arriving (Next Week):**\n# {list(e['incoming_shipments'])[0] if e['incoming_shipments'] else 0}")
            
        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] 
            arrived = 0
            if name == "Factory": 
                 if state['factory_production_pipeline']: arrived = list(state['factory_production_pipeline'])[0]
            else: 
                 if e_curr['incoming_shipments']: arrived = list(e_curr['incoming_shipments'])[0]
                        
            inv_after_arrival = e_curr['inventory'] + arrived
                        
            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, 
                '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.")
                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"):
                    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*.
            
            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:** 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).")
                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"):
                    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)
                                        
                    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'):
                st.write("Your weekly history will be displayed here after you complete the first week.")
            else:
                try:
                    history_df = pd.json_normalize(state['logs'])
                    human_cols = {
                        'week': 'Week', f'{human_role}.opening_inventory': 'Opening Inv.',
                        f'{human_role}.opening_backlog': 'Opening Backlog', f'{human_role}.arrived_this_week': 'Arrived This Week',
                        f'{human_role}.incoming_order': 'Incoming Order', f'{human_role}.initial_order': 'Your Initial Order',
                        f'{human_role}.ai_suggestion': 'AI Suggestion', f'{human_role}.order_placed': 'Your Final Order',
                        f'{human_role}.arriving_next_week': 'Arriving Next Week', f'{human_role}.weekly_cost': 'Weekly Cost',
                    }
                    ordered_display_cols_keys = [
                        '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' 
                     ]
                    final_cols_to_display = [col for col in ordered_display_cols_keys if col in history_df.columns]
                    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)
                        if 'Weekly Cost' in display_df.columns:
                            display_df['Weekly Cost'] = display_df['Weekly Cost'].apply(lambda x: f"${x:,.2f}" if isinstance(x, (int, float)) 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")
        st.sidebar.markdown(f"**Game ID**: `{state['participant_id']}`\n\n**Current Week**: {week}")
        if st.sidebar.button("🔄 Reset Game"):
            if 'final_order_input' in st.session_state: del st.session_state.final_order_input
            del st.session_state.game_state
            st.rerun()

    # --- Game Over Interface ---
    if 'game_state' in st.session_state and not st.session_state.game_state.get('game_running', False) and st.session_state.game_state['week'] > WEEKS:
        st.header("🎉 Game Over!")
        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 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()