Ride-Share Surge Control

from tys import probe, progress


def simulate(cfg: dict):
    import math
    import random
    from collections import deque

    import simpy

    env = simpy.Environment()

    horizon = int(cfg.get("time_horizon", 240))
    seed = int(cfg.get("seed", 0))
    rng = random.Random(seed)

    base_demand = float(cfg["base_demand"])
    demand_elasticity = float(cfg["demand_elasticity"])
    supply_elasticity = float(cfg["supply_elasticity"])
    trip_time = float(cfg["trip_time"])
    baseline_active = float(cfg["baseline_active_drivers"])

    target_wait = float(cfg.get("target_wait", 3.0))
    wait_window = int(cfg.get("wait_window", 10))

    price_gain = float(cfg.get("price_gain", 0.0))
    price_cap = float(cfg.get("price_cap", 3.0))
    smoothing_alpha = float(cfg.get("smoothing_alpha", 0.0))

    initial_price = float(cfg.get("initial_price", 1.0))
    price_delay = int(cfg.get("price_delay_minutes", 0))
    supply_tau = float(cfg.get("supply_time_constant", 1.0))

    shock_start = int(cfg.get("shock_start", 0))
    shock_end = int(cfg.get("shock_end", 0))
    shock_multiplier = float(cfg.get("shock_multiplier", 1.0))

    assert wait_window > 0, "wait_window must be positive"
    assert supply_tau > 0, "supply_time_constant must be positive"
    assert price_cap >= 1.0, "price_cap must be at least 1.0"

    def poisson(lam: float) -> int:
        """Sample a Poisson random variate with mean ``lam``."""
        if lam <= 0:
            return 0
        L = math.exp(-lam)
        k = 0
        p = 1.0
        while p > L:
            k += 1
            p *= rng.random()
        return k - 1

    price = max(1.0, initial_price)
    price_history = deque([price] * (max(price_delay, 0) + 1), maxlen=max(price_delay, 0) + 1)

    queue = 0.0
    idle_drivers = baseline_active
    engaged_drivers = 0.0

    match_history = deque([base_demand] * wait_window, maxlen=wait_window)

    total_matches = 0.0
    wait_samples = []
    price_samples = []

    def recorder():
        while True:
            probe("riders_waiting", env.now, queue)
            probe("idle_drivers", env.now, idle_drivers)
            probe("drivers_on_trip", env.now, engaged_drivers)
            yield env.timeout(1)

    env.process(recorder())
    done = env.event()

    def dynamics():
        nonlocal price, queue, idle_drivers, engaged_drivers, total_matches

        for minute in range(horizon):
            shock = shock_multiplier if shock_start <= minute < shock_end else 1.0
            arrival_rate = max(0.0, base_demand * (price ** demand_elasticity) * shock)
            arrivals = poisson(arrival_rate)
            queue += arrivals

            price_for_supply = price_history[0]
            target_active = baseline_active * (price_for_supply ** supply_elasticity)
            active_total = idle_drivers + engaged_drivers
            delta_active = (target_active - active_total) / supply_tau
            proposed_active = max(active_total + delta_active, 0.0)
            change_active = proposed_active - active_total

            if change_active >= 0:
                driver_logins = change_active
                driver_logouts = 0.0
                idle_drivers += driver_logins
            else:
                driver_logins = 0.0
                driver_logouts = min(-change_active, idle_drivers)
                idle_drivers -= driver_logouts
                proposed_active = idle_drivers + engaged_drivers
                change_active = proposed_active - active_total

            capacity = min(idle_drivers, queue)
            matches = max(0.0, capacity)
            queue -= matches
            idle_drivers -= matches
            engaged_drivers += matches
            total_matches += matches

            completions = engaged_drivers / trip_time if trip_time > 0 else engaged_drivers
            completions = min(completions, engaged_drivers)
            engaged_drivers -= completions
            idle_drivers += completions

            match_history.append(max(matches, 1e-6))
            avg_service = max(sum(match_history) / len(match_history), 1e-6)
            wait_estimate = queue / avg_service

            error = (wait_estimate - target_wait) / max(target_wait, 1e-6)
            price_raw = price + price_gain * error
            price_raw = min(max(price_raw, 1.0), price_cap)
            if 0.0 < smoothing_alpha < 1.0:
                price_next = (1 - smoothing_alpha) * price + smoothing_alpha * price_raw
            else:
                price_next = price_raw

            price_history.append(price_next)

            probe("price_multiplier", env.now, price)
            probe("wait_estimate", env.now, wait_estimate)
            probe("arrival_rate", env.now, arrival_rate)
            probe("driver_logins", env.now, driver_logins)
            probe("driver_logouts", env.now, driver_logouts)
            probe("matches_per_min", env.now, matches)
            probe("trip_completions", env.now, completions)

            price_samples.append(price)
            wait_samples.append(wait_estimate)

            if minute == shock_start:
                progress(int(100 * minute / horizon), "🌧️ Demand shock hits")
            if minute == shock_end:
                progress(int(100 * minute / horizon), "☀️ Demand normalises")
            if price_next >= price_cap * 0.95:
                progress(int(100 * minute / horizon), "⚠️ Surge near cap")
            if wait_estimate > 2 * target_wait:
                progress(int(100 * minute / horizon), "⏱️ Waits running hot")

            price = price_next
            yield env.timeout(1)

        progress(100)

        sorted_waits = sorted(wait_samples)

        def percentile(data, pct):
            if not data:
                return 0.0
            k = (len(data) - 1) * pct
            f = math.floor(k)
            c = math.ceil(k)
            if f == c:
                return data[int(k)]
            return data[f] * (c - k) + data[c] * (k - f)

        avg_wait = sum(wait_samples) / len(wait_samples) if wait_samples else 0.0
        avg_price = sum(price_samples) / len(price_samples) if price_samples else price
        price_variance = sum((p - avg_price) ** 2 for p in price_samples) / len(price_samples) if price_samples else 0.0

        done.succeed({
            "avg_wait_minutes": avg_wait,
            "p95_wait_minutes": percentile(sorted_waits, 0.95),
            "peak_wait_minutes": max(wait_samples) if wait_samples else 0.0,
            "avg_price_multiplier": avg_price,
            "price_volatility": math.sqrt(price_variance),
            "completed_trips": total_matches,
        })

    env.process(dynamics())
    env.run(until=done)
    return done.value


def requirements():
    """Dependencies needed for this simulation."""
    return {
        "external": ["simpy==4.1.1"],
    }