simulator.py

"""
simulator.py — Venue seeding, attendee state machine, crowd density tracking.
All state lives in module-level dicts; no database.
"""

import random
import asyncio
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional

# ---------------------------------------------------------------------------
# Venue definition — Wankhede Stadium, Mumbai
# ---------------------------------------------------------------------------

VENUE_NAME = "Wankhede Stadium"
EVENT_NAME = "IPL Final 2025 — Mumbai Indians vs Chennai Super Kings"
VENUE_CENTER = {"lat": 18.9388, "lng": 72.8252}

# 8 gates placed around the stadium perimeter (~150 m radius)
GATES: Dict[str, dict] = {
    "G1": {
        "name": "Gate 1 — North Stand",
        "lat": 18.9403,
        "lng": 72.8252,
        "capacity": 2200,
        "direction": "North",
        "transit": "metro",
    },
    "G2": {
        "name": "Gate 2 — NE Corner",
        "lat": 18.9399,
        "lng": 72.8268,
        "capacity": 1600,
        "direction": "NE",
        "transit": "metro",
    },
    "G3": {
        "name": "Gate 3 — East Stand",
        "lat": 18.9388,
        "lng": 72.8272,
        "capacity": 2000,
        "direction": "East",
        "transit": "metro",
    },
    "G4": {
        "name": "Gate 4 — SE Corner",
        "lat": 18.9377,
        "lng": 72.8268,
        "capacity": 1500,
        "direction": "SE",
        "transit": "bus",
    },
    "G5": {
        "name": "Gate 5 — South Stand",
        "lat": 18.9373,
        "lng": 72.8252,
        "capacity": 2400,
        "direction": "South",
        "transit": "bus",
    },
    "G6": {
        "name": "Gate 6 — SW Corner",
        "lat": 18.9377,
        "lng": 72.8236,
        "capacity": 1400,
        "direction": "SW",
        "transit": "bus",
    },
    "G7": {
        "name": "Gate 7 — West Stand",
        "lat": 18.9388,
        "lng": 72.8232,
        "capacity": 2000,
        "direction": "West",
        "transit": "taxi",
    },
    "G8": {
        "name": "Gate 8 — NW Corner",
        "lat": 18.9399,
        "lng": 72.8236,
        "capacity": 1700,
        "direction": "NW",
        "transit": "taxi",
    },
}

TRANSIT_NODES: Dict[str, dict] = {
    "metro": {
        "name": "Churchgate Metro Station",
        "lat": 18.9352,
        "lng": 72.8259,
        "capacity": 6000,
        "walk_min": 8,
        "icon": "🚇",
    },
    "bus": {
        "name": "Marine Lines Bus Depot",
        "lat": 18.9330,
        "lng": 72.8245,
        "capacity": 4000,
        "walk_min": 12,
        "icon": "🚌",
    },
    "taxi": {
        "name": "Wankhede Taxi Stand",
        "lat": 18.9410,
        "lng": 72.8268,
        "capacity": 2500,
        "walk_min": 5,
        "icon": "🚕",
    },
}

# 20 sections A–T with gate assignments
SECTIONS: Dict[str, dict] = {
    "A": {"gate": "G1", "size": 750, "stand": "North"},
    "B": {"gate": "G1", "size": 750, "stand": "North"},
    "C": {"gate": "G1", "size": 750, "stand": "North"},
    "D": {"gate": "G2", "size": 750, "stand": "North-East"},
    "E": {"gate": "G2", "size": 750, "stand": "North-East"},
    "F": {"gate": "G3", "size": 750, "stand": "East"},
    "G": {"gate": "G3", "size": 750, "stand": "East"},
    "H": {"gate": "G3", "size": 750, "stand": "East"},
    "I": {"gate": "G4", "size": 750, "stand": "South-East"},
    "J": {"gate": "G4", "size": 750, "stand": "South-East"},
    "K": {"gate": "G5", "size": 750, "stand": "South"},
    "L": {"gate": "G5", "size": 750, "stand": "South"},
    "M": {"gate": "G5", "size": 750, "stand": "South"},
    "N": {"gate": "G6", "size": 750, "stand": "South-West"},
    "O": {"gate": "G6", "size": 750, "stand": "South-West"},
    "P": {"gate": "G7", "size": 750, "stand": "West"},
    "Q": {"gate": "G7", "size": 750, "stand": "West"},
    "R": {"gate": "G7", "size": 750, "stand": "West"},
    "S": {"gate": "G8", "size": 750, "stand": "North-West"},
    "T": {"gate": "G8", "size": 750, "stand": "North-West"},
}

# Gate → departure wave (0–7). Adjacent sections staggered so no pile-up at T+0.
GATE_WAVE: Dict[str, int] = {
    "G1": 0, "G2": 3, "G3": 1, "G4": 4,
    "G5": 2, "G6": 5, "G7": 3, "G8": 6,
}

# Wave start offsets in minutes from event end
WAVE_OFFSET_MIN = [0, 6, 12, 18, 24, 30, 36, 42]

# ---------------------------------------------------------------------------
# Attendee dataclass
# ---------------------------------------------------------------------------

@dataclass
class Attendee:
    ticket_id: str
    section: str
    gate_id: str
    seat_row: int
    wave: int
    transit: str
    depart_offset_min: float   # minutes after event end to leave seat
    state: str = "seated"      # seated | walking | at_gate | dispersed
    state_entered_at: float = field(default_factory=time.monotonic)

# ---------------------------------------------------------------------------
# Shared simulation state (module-level, in-memory)
# ---------------------------------------------------------------------------

attendees: Dict[str, Attendee] = {}
gate_state: Dict[str, dict] = {}
transit_state: Dict[str, dict] = {}
dispersal_stats: dict = {}
active_alerts: List[dict] = []
dispersal_strategy_text: str = ""   # Cached Gemini output
sim_start_time: float = 0.0
sim_tick: int = 0

# Track how long each gate has been dense (for alert rule)
_gate_dense_ticks: Dict[str, int] = {}
_alert_seen: Dict[str, float] = {}  # gate_id → last alert timestamp

# ---------------------------------------------------------------------------
# Seeding
# ---------------------------------------------------------------------------

def seed(strategy_override: Optional[dict] = None) -> None:
    """Seed 15,000 attendees and initialise crowd state."""
    global sim_start_time, sim_tick

    random.seed(42)
    attendees.clear()
    gate_state.clear()
    transit_state.clear()
    active_alerts.clear()
    _gate_dense_ticks.clear()
    _alert_seen.clear()

    sim_start_time = time.monotonic()
    sim_tick = 0

    # Initialise gate crowd state
    for gid, gdata in GATES.items():
        gate_state[gid] = {
            "gate_id": gid,
            "name": gdata["name"],
            "lat": gdata["lat"],
            "lng": gdata["lng"],
            "capacity": gdata["capacity"],
            "current_count": 0,
            "density_pct": 0.0,
            "color": "green",
            "transit": gdata["transit"],
        }
        _gate_dense_ticks[gid] = 0

    # Initialise transit crowd state
    for tid, tdata in TRANSIT_NODES.items():
        transit_state[tid] = {
            "transit_id": tid,
            "name": tdata["name"],
            "lat": tdata["lat"],
            "lng": tdata["lng"],
            "capacity": tdata["capacity"],
            "current_count": 0,
            "density_pct": 0.0,
            "color": "green",
            "walk_min": tdata["walk_min"],
        }

    # Seed attendees
    ticket_counter = 1
    for section, sdata in SECTIONS.items():
        gate_id = sdata["gate"]
        transit_id = GATES[gate_id]["transit"]
        base_wave = GATE_WAVE[gate_id]

        # Allow strategy override to adjust wave
        if strategy_override and gate_id in strategy_override:
            base_wave = strategy_override[gate_id].get("wave", base_wave)

        for i in range(sdata["size"]):
            ticket_id = f"TKT-{ticket_counter:05d}"
            ticket_counter += 1

            wave = base_wave
            # Small jitter within wave (±2 min) to avoid hard spikes
            offset = WAVE_OFFSET_MIN[wave] + random.uniform(-2, 2)
            offset = max(0, offset)

            attendees[ticket_id] = Attendee(
                ticket_id=ticket_id,
                section=section,
                gate_id=gate_id,
                seat_row=random.randint(1, 30),
                wave=wave,
                transit=transit_id,
                depart_offset_min=offset,
                state="seated",
                state_entered_at=sim_start_time,
            )

    _update_stats()


def get_demo_ticket() -> str:
    """Return a consistent demo ticket ID for the landing page."""
    return "TKT-00042"


def get_attendee_context(ticket_id: str) -> Optional[dict]:
    """Return a rich dict describing the attendee's plan (for Gemini system prompt)."""
    a = attendees.get(ticket_id)
    if not a:
        return None

    gate = GATES[a.gate_id]
    transit = TRANSIT_NODES[a.transit]
    wave_offset = WAVE_OFFSET_MIN[a.wave]

    return {
        "ticket_id": ticket_id,
        "section": a.section,
        "seat_row": a.seat_row,
        "gate_id": a.gate_id,
        "gate_name": gate["name"],
        "gate_direction": gate["direction"],
        "transit_name": transit["name"],
        "transit_icon": transit["icon"],
        "transit_walk_min": transit["walk_min"],
        "wave": a.wave,
        "depart_offset_min": round(a.depart_offset_min, 1),
        "state": a.state,
        "transit_id": a.transit,
        "event_name": EVENT_NAME,
        "venue_name": VENUE_NAME,
        "gate_density_pct": gate_state[a.gate_id]["density_pct"],
        "gate_color": gate_state[a.gate_id]["color"],
    }


# ---------------------------------------------------------------------------
# Simulation tick
# ---------------------------------------------------------------------------

SECONDS_PER_SIM_MINUTE = 3   # 1 sim-minute = 3 real seconds → 45 sim-min = 135 s

def _elapsed_sim_minutes() -> float:
    elapsed_real_sec = time.monotonic() - sim_start_time
    return elapsed_real_sec / SECONDS_PER_SIM_MINUTE


def tick() -> None:
    """Advance simulation state. Called every 5 real seconds."""
    global sim_tick
    sim_tick += 1
    elapsed = _elapsed_sim_minutes()

    # Reset gate/transit counts
    for gs in gate_state.values():
        gs["current_count"] = 0
    for ts in transit_state.values():
        ts["current_count"] = 0

    for a in attendees.values():
        if a.state == "seated":
            if elapsed >= a.depart_offset_min:
                a.state = "walking"
                a.state_entered_at = time.monotonic()

        elif a.state == "walking":
            # Walking takes 2–4 sim minutes
            walk_time = (time.monotonic() - a.state_entered_at) / SECONDS_PER_SIM_MINUTE
            if walk_time >= random.uniform(2, 4):
                a.state = "at_gate"
                a.state_entered_at = time.monotonic()
                gate_state[a.gate_id]["current_count"] += 1

        elif a.state == "at_gate":
            gate_state[a.gate_id]["current_count"] += 1
            # Processing at gate takes 3–7 sim minutes
            gate_time = (time.monotonic() - a.state_entered_at) / SECONDS_PER_SIM_MINUTE
            if gate_time >= random.uniform(3, 7):
                a.state = "dispersed"
                a.state_entered_at = time.monotonic()
                transit_state[a.transit]["current_count"] += 1

        elif a.state == "dispersed":
            # Still contribute to transit count briefly (5 sim-min)
            disp_time = (time.monotonic() - a.state_entered_at) / SECONDS_PER_SIM_MINUTE
            if disp_time < 5:
                transit_state[a.transit]["current_count"] += 1

    # Update density %
    for gid, gs in gate_state.items():
        raw = gs["current_count"] / gs["capacity"]
        # Add Gaussian noise for realism
        noisy = max(0.0, min(1.0, raw + random.gauss(0, 0.02)))
        gs["density_pct"] = round(noisy, 3)
        gs["color"] = _density_color(noisy)

        # Track dense ticks for alert rule
        if noisy > 0.85:
            _gate_dense_ticks[gid] += 1
        else:
            _gate_dense_ticks[gid] = 0

    for tid, ts in transit_state.items():
        raw = ts["current_count"] / ts["capacity"]
        noisy = max(0.0, min(1.0, raw + random.gauss(0, 0.01)))
        ts["density_pct"] = round(noisy, 3)
        ts["color"] = _density_color(noisy)

    _update_alerts()
    _update_stats()


def _density_color(density: float) -> str:
    if density < 0.60:
        return "green"
    if density < 0.85:
        return "yellow"
    return "red"


def _update_alerts() -> None:
    """Apply the >85% for >2 ticks rule. Does NOT call Gemini (done in main.py)."""
    seen_gates = {a["gate_id"] for a in active_alerts}

    for gid, dense_ticks in _gate_dense_ticks.items():
        if dense_ticks >= 2 and gid not in seen_gates:
            gs = gate_state[gid]
            # Find affected sections
            affected = [s for s, sd in SECTIONS.items() if sd["gate"] == gid]
            pending = sum(
                1 for a in attendees.values()
                if a.gate_id == gid and a.state in ("seated", "walking", "at_gate")
            )
            active_alerts.append({
                "alert_id": f"ALT-{gid}-{sim_tick}",
                "gate_id": gid,
                "gate_name": gs["name"],
                "density_pct": gs["density_pct"],
                "dense_ticks": dense_ticks,
                "affected_sections": affected,
                "pending_attendees": pending,
                "suggestion": None,   # Filled by gemini_client in main.py
                "created_at": time.time(),
            })

    # Remove stale alerts where gate is now green or resolved
    to_remove = []
    for alert in active_alerts:
        gid = alert["gate_id"]
        if gate_state[gid]["color"] == "green" and (time.time() - alert["created_at"]) > 60:
            to_remove.append(alert)
    for a in to_remove:
        active_alerts.remove(a)


def _update_stats() -> None:
    total = len(attendees)
    seated = sum(1 for a in attendees.values() if a.state == "seated")
    walking = sum(1 for a in attendees.values() if a.state == "walking")
    at_gate = sum(1 for a in attendees.values() if a.state == "at_gate")
    dispersed = sum(1 for a in attendees.values() if a.state == "dispersed")

    dispersal_stats["total"] = total
    dispersal_stats["seated"] = seated
    dispersal_stats["walking"] = walking
    dispersal_stats["at_gate"] = at_gate
    dispersal_stats["dispersed"] = dispersed
    dispersal_stats["pct_dispersed"] = round(dispersed / total * 100, 1) if total else 0
    dispersal_stats["elapsed_sim_min"] = round(_elapsed_sim_minutes(), 1)

    # Average wait estimate: remaining attendees × factor
    remaining = total - dispersed
    dispersal_stats["avg_wait_min"] = round((remaining / total) * 22, 1) if total else 0


def get_crowd_state_snapshot() -> dict:
    """Serialisable snapshot for /api/crowd_state."""
    return {
        "tick": sim_tick,
        "elapsed_sim_min": dispersal_stats.get("elapsed_sim_min", 0),
        "gates": list(gate_state.values()),
        "transit": list(transit_state.values()),
        "stats": dict(dispersal_stats),
    }