L03: ATL Checkpoint

"""Figures and the animated checkpoint for L03 (called by sim.main)."""
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.collections import PatchCollection
from matplotlib.lines import Line2D
from matplotlib.patches import Patch, Rectangle
from matplotlib.ticker import MaxNLocator

import sim
from sim import H, LANE_MAX, OUT, STREAMS, T_END, T_OPEN, TARGET_MIN, TDC_PODIUMS

# Dark theme (validated categorical steps; blue/orange/aqua pass all-pairs CVD checks)
BG, PANEL, INK, INK2, MUTED, GRID, AXIS = "#0d0d0d", "#1a1a19", "#ffffff", "#c3c2b7", "#898781", "#2c2c2a", "#383835"
BLUE, ORANGE, AQUA, YELLOW, RED = "#3987e5", "#d95926", "#199e70", "#c98500", "#d03b3b"
STREAM_C = (BLUE, AQUA)
POLICY = {"SIPP (Erlang C)": (ORANGE, "-"), "SimOpt naive": (AQUA, "--"),
          "SimOpt robust": (BLUE, "-"), "Flat": (MUTED, ":")}
plt.rcParams.update({
    "figure.facecolor": BG, "savefig.facecolor": BG, "axes.facecolor": PANEL, "axes.edgecolor": AXIS,
    "axes.labelcolor": INK2, "text.color": INK, "xtick.color": MUTED, "ytick.color": MUTED,
    "axes.grid": True, "axes.axisbelow": True, "grid.color": GRID, "grid.linewidth": 0.6,
    "axes.spines.top": False, "axes.spines.right": False, "font.size": 10, "axes.titlesize": 12,
    "axes.titleweight": "bold", "axes.titlelocation": "left", "legend.frameon": False, "lines.linewidth": 2,
})


def hour_axis(ax, lo=4, hi=24, step=2):
    ax.set_xlim(lo, hi)
    ax.set_xticks(range(lo, hi + 1, step), [f"{h:02d}:00" for h in range(lo, hi + 1, step)])


def save(fig, name):
    fig.savefig(OUT / name, dpi=130, bbox_inches="tight")
    plt.close(fig)


def steps(plan_row):
    """x, y for a step plot of an hourly plan."""
    return np.arange(H + 1) + T_OPEN / 60, np.append(plan_row, plan_row[-1])


def plural(n, word):
    return f"{n} {word}" + ("" if n == 1 else "s")


# ----------------------------------------------------------------------------------------
def fig_nhpp(n):
    fig, (a, b) = plt.subplots(1, 2, figsize=(13, 4.6), gridspec_kw=dict(width_ratios=[1, 1.4]))
    lo, hi = 5 * 60, 6 * 60
    m = (n["cand"] >= lo) & (n["cand"] < hi)
    c, y, k = n["cand"][m] / 60, n["y"][m], n["keep"][m]
    a.scatter(c[~k], y[~k], s=5, color=MUTED, alpha=0.45, lw=0, label=f"rejected ({(~k).sum()})")
    a.scatter(c[k], y[k], s=5, color=BLUE, lw=0, label=f"accepted = arrivals ({k.sum()})")
    t = np.linspace(lo, hi, 400)
    a.plot(t / 60, sim.rate(t), color=INK, lw=2, label=r"$\lambda(t)$")
    a.axhline(sim.LAM_MAX, color=ORANGE, ls="--", lw=1.5, label=r"$\lambda_{max}$ (candidate rate)")
    a.set_xlim(5, 6)
    a.set_xticks([5, 5.25, 5.5, 5.75, 6], ["05:00", "05:15", "05:30", "05:45", "06:00"])
    a.set_ylim(0, sim.LAM_MAX * 1.18)
    a.set(ylabel="rate (passengers / min)", title="Thinning: keep points under lambda(t)")
    a.legend(loc="upper right", ncol=2, fontsize=8.5)

    ctr = (n["bins"][:-1] + 7.5) / 60
    mean, sd = n["counts"].mean(0), n["counts"].std(0, ddof=1)
    b.fill_between(ctr, mean - 2 * sd, mean + 2 * sd, color=BLUE, alpha=0.3, lw=0, label="simulated, +/-2 sd (day-to-day)")
    b.plot(ctr, mean, color=BLUE, label=f"simulated mean of {len(n['totals'])} days")
    b.plot(ctr, n["expected"], color=ORANGE, ls="--", lw=1.6, label=r"expected $\Lambda$(bin) = $\int\lambda\,dt$")
    hour_axis(b, 4, 22)
    b.set_ylim(0, (mean + 2 * sd).max() * 1.35)
    b.set(ylabel="arrivals per 15 min", title="NHPP arrivals at one checkpoint, 22,000/day")
    tot = n["totals"]
    b.text(0.99, 0.70, f"daily total: mean {tot.mean():,.0f} vs Lambda {n['Lambda']:,.0f}\n"
           f"variance / mean = {tot.var(ddof=1) / tot.mean():.2f} (Poisson: 1)",
           transform=b.transAxes, ha="right", va="top", color=INK2, fontsize=9)
    b.legend(loc="upper right", fontsize=8.5)
    save(fig, "nhpp_thinning.png")


def fig_erlang(m):
    fig, (a, b) = plt.subplots(1, 2, figsize=(13, 4.6))
    rho = np.linspace(0.4, 0.965, 200)
    lam = rho * sim.MMC_C * sim.MMC_MU
    a.plot(rho, [sim.mmc_wq(l, sim.MMC_MU, sim.MMC_C) for l in lam], color=ORANGE, lw=2, label="Erlang C (exact)")
    sw = m["sweep"]
    a.errorbar(sw[:, 0], sw[:, 1], yerr=sw[:, 2], fmt="o", ms=7, color=BLUE, ecolor=BLUE, capsize=5, lw=2,
               mec=PANEL, mew=2, label=f"simulation, mean +/- 95% CI ({m['R']} reps)")
    a.set_yscale("log")
    a.set(xlabel=r"utilisation $\rho = \lambda / (c\mu)$", ylabel="mean wait in queue Wq (min)",
          title=f"M/M/{sim.MMC_C}: simulated Wq vs queueing theory")
    a.legend(loc="upper left")
    a.text(0.98, 0.04, f"SimPy (process interaction) and heapq (event scheduling)\n"
           f"give identical waits on the same inputs: max |dW| = {m['engines_maxdiff']:.0e} min",
           transform=a.transAxes, ha="right", va="bottom", fontsize=8.5, color=INK2)

    t = m["surv_t"]
    mean, hw = sim.t_ci(m["surv"])
    b.fill_between(t, np.clip(mean - hw, 1e-6, None), mean + hw, color=BLUE, alpha=0.3, lw=0,
                   label=f"SimPy, 95% CI across {m['R']} replications")
    b.plot(t, mean, color=BLUE, lw=2)
    b.plot(t, sim.mmc_p_wait_gt(t, sim.MMC_LAM, sim.MMC_MU, sim.MMC_C), color=ORANGE, ls="--", lw=2,
           label=r"Erlang C: $C(c,a)\,e^{-(c\mu-\lambda)t}$")
    b.set_yscale("log")
    b.set_ylim(1e-5, 1.2)
    b.set(xlabel="t (min)", ylabel="P(Wq > t)",
          title=f"Whole wait distribution at rho = {sim.MMC_LAM / sim.MMC_C / sim.MMC_MU:.2f}")
    b.legend()
    save(fig, "verification_erlangc.png")


def fig_warmup(wu, wq):
    fig, (a, b) = plt.subplots(1, 2, figsize=(13, 4.6))
    i = np.arange(len(wu["ybar"]))
    a.plot(i, wu["ybar"], color=MUTED, lw=0.7, alpha=0.7, label="customer i's wait, mean over 500 reps")
    a.plot(np.arange(len(wu["ma"])), wu["ma"], color=BLUE, lw=2.5, label=f"Welch moving average (w = {wu['w']})")
    a.axhline(wq, color=ORANGE, ls="--", lw=1.6, label="steady state (Erlang C)")
    a.axvspan(0, wu["d"], color=RED, alpha=0.15, lw=0)
    a.text(wu["d"] + 40, 0.12 * wq, f"delete first {wu['d']}", color=INK2, fontsize=9)
    a.set(xlabel="customer index i (queue starts empty)", ylabel="wait (min)", xlim=(0, len(i)),
          title="Welch's procedure: where does the start-up bias die out?")
    a.set_ylim(0, None)
    a.legend(loc="lower right", fontsize=8.5)

    ns, cov = np.array(sim.COVER_NS), wu["coverage"]
    se = 1.96 * np.sqrt(0.95 * 0.05 / wu["macro"])
    b.axhspan(0.95 - se, 0.95 + se, color=MUTED, alpha=0.15, lw=0)
    b.axhline(0.95, color=MUTED, ls="--", lw=1.2)
    b.text(ns[0], 0.95 + se + 0.004, "nominal 95% (+/- Monte Carlo error)", color=INK2, fontsize=8.5)
    b.plot(ns, cov[:, 0], "o-", color=ORANGE, ms=8, mec=PANEL, mew=2, label="no deletion (biased low)")
    b.plot(ns, cov[:, 1], "o-", color=BLUE, ms=8, mec=PANEL, mew=2, label=f"delete first {wu['d']} (Welch)")
    b.set_xscale("log", base=2)
    b.set_xticks(ns, [f"{n:,}" for n in ns])
    b.set(xlabel="customers per replication", ylabel="fraction of CIs covering true Wq",
          title=f"Payoff: coverage of 95% t-CIs ({wu['macro']} experiments x 10 reps)")
    b.legend(loc="lower right")
    save(fig, "warmup_coverage.png")


def fig_staffing(plans, summ, traces):
    fig = plt.figure(figsize=(13, 10.8))
    gs = fig.add_gridspec(3, 2, height_ratios=[1, 1, 0.8], hspace=0.4, wspace=0.14, top=0.9)
    lam_h = sim.expected_per_hour()
    for k in (0, 1):
        a = fig.add_subplot(gs[0, k])
        share = sim.P_PRECHECK if k else 1 - sim.P_PRECHECK
        load = lam_h * share / 60 * sim.LANE_MEAN[k]  # offered load = lanes busy on average
        a.fill_between(*steps(load), step="post", color=MUTED, alpha=0.25, lw=0)
        for name, p in plans.items():
            c, ls = POLICY[name]
            a.step(*steps(p[k]), where="post", color=c, ls=ls, lw=2 if name != "Flat" else 1.5)
        hour_axis(a)
        a.yaxis.set_major_locator(MaxNLocator(integer=True))
        a.set(ylabel="lanes open", title=f"{STREAMS[k]} lanes by hour")
        a.set_ylim(0, None)

        b = fig.add_subplot(gs[1, k])
        x = np.arange(H) + T_OPEN / 60 + 0.5
        for name, s in summ.items():
            c, ls = POLICY[name]
            b.plot(x, s["hourly"][k], color=c, ls=ls, marker="o", ms=4, lw=1.8 if name != "Flat" else 1.2)
        b.axhline(TARGET_MIN, color=RED, lw=1.5)
        hour_axis(b)
        b.set_ylim(0, 40 if k == 0 else 25)
        b.set(ylabel="P95 wait (min) of that hour's arrivals",
              title=f"{STREAMS[k]}: out-of-sample check (40 fresh days)")
        if k == 0:
            b.text(4.4, 38.5, f"Flat leaves the chart (worst hour {summ['Flat']['worst0']:.0f} min)",
                   va="top", color=MUTED, fontsize=8.5)
    handles = [Patch(color=MUTED, alpha=0.4, label="offered load lambda/mu (lanes busy on average)")]
    handles += [Line2D([], [], color=POLICY[n][0], ls=POLICY[n][1], lw=2, label=f"{n}: {p.sum()} lane-h/day")
                for n, p in plans.items()]
    handles.append(Line2D([], [], color=RED, lw=1.5, label="target: P95 wait <= 15 min"))
    fig.legend(handles=handles, loc="upper center", ncol=3, bbox_to_anchor=(0.5, 0.975), fontsize=9.5)

    c = fig.add_subplot(gs[2, :])
    sipp_cost = plans["SIPP (Erlang C)"].sum()
    for key, name in (("naive", "SimOpt naive"), ("robust", "SimOpt robust")):
        tr = traces[key]
        col, ls = POLICY[name]
        c.step(tr[:, 0], tr[:, 1], where="post", color=col, ls=ls, label=f"{name}: {len(tr)} evaluations")
    c.axhline(sipp_cost, color=ORANGE, lw=1.2, ls=":")
    c.text(max(len(t) for t in traces.values()), sipp_cost - 0.4, f"SIPP start: {sipp_cost} lane-h",
           color=ORANGE, fontsize=9, ha="right", va="top")
    c.set(xlabel="simulation evaluations (naive: 12 replicated days each, robust: 24; same seeds = CRN)",
          ylabel="lane-hours of plan", title="Greedy forward sweep: hour by hour, drop lanes until the constraint binds")
    c.legend(loc="lower left")
    save(fig, "staffing.png")


def fig_policies(summ, reps):
    fig, (a, b) = plt.subplots(1, 2, figsize=(13, 4.8))
    offset = {"SIPP (Erlang C)": (10, 0, "center"), "SimOpt naive": (8, 6, "bottom"),
              "SimOpt robust": (8, -8, "top"), "Flat": (10, 0, "center")}  # keep labels apart
    for name, s in summ.items():
        col, _ = POLICY[name]
        dx, dy, va = offset[name]
        a.scatter(s["lane_hours"], s["worst0"], s=110, color=col, edgecolor=PANEL, lw=2, zorder=3)
        a.annotate(f"{name}\n{plural(s['viol0'] + s['viol1'], 'hour')} over 15 min",
                   (s["lane_hours"], s["worst0"]), xytext=(dx, dy), textcoords="offset points",
                   color=INK2, fontsize=9, va=va)
    a.axhline(TARGET_MIN, color=RED, lw=1.5)
    a.axhspan(1, TARGET_MIN, color=AQUA, alpha=0.07, lw=0)
    a.set_yscale("log")
    a.set_ylim(5, 300)
    a.set_xlim(min(s["lane_hours"] for s in summ.values()) - 4, max(s["lane_hours"] for s in summ.values()) + 14)
    a.set(xlabel="cost: lane-hours per day", ylabel="worst-hour P95 wait, Standard (min)",
          title="Cost vs service (validated on 40 fresh days)")
    a.text(a.get_xlim()[0] + 0.5, TARGET_MIN * 0.93, "feasible", color=AQUA, fontsize=9, va="top")

    for name, rr in reps.items():
        col, ls = POLICY[name]
        w = np.sort(np.concatenate([r["wait"][r["kind"] == 0] for r in rr]))
        b.plot(w, np.arange(1, len(w) + 1) / len(w), color=col, ls=ls, label=name)
    b.axvline(TARGET_MIN, color=RED, lw=1.5)
    b.axhline(0.95, color=MUTED, lw=1, ls="--")
    b.set(xlim=(0, 60), ylim=(0, 1.01), xlabel="wait before screening (min)", ylabel="fraction of passengers",
          title="Standard stream: wait-time CDF, all day")
    b.legend(loc="lower right")
    save(fig, "policy_comparison.png")


def mean_band(ax, nq_reps, color, label, spaghetti=False):
    x = sim.MINUTES / 60
    if spaghetti:
        for q in nq_reps:
            ax.plot(x, q, color=color, lw=0.5, alpha=0.10)
    m, hw = sim.t_ci(nq_reps)
    ax.fill_between(x, m - hw, m + hw, color=color, alpha=0.35, lw=0)
    ax.plot(x, m, color=color, lw=2, label=label)


def fig_queue(reps):
    fig, axes = plt.subplots(2, 1, figsize=(13, 7), sharex=True)
    for k, ax in enumerate(axes):
        for name in ("SIPP (Erlang C)", "SimOpt robust"):
            mean_band(ax, np.array([r["nq"][k] for r in reps[name]]), POLICY[name][0], f"{name}: mean, 95% CI band")
        ax.set(ylabel="passengers waiting", title=f"{STREAMS[k]} queue over the day (40 replications)")
        ax.set_ylim(0, None)
        ax.legend(loc="upper right")
    hour_axis(axes[1])
    save(fig, "queue_over_day.png")


def fig_crn(c):
    fig, (a, b) = plt.subplots(1, 2, figsize=(13, 4.8), gridspec_kw=dict(width_ratios=[1, 1.15]))
    a.scatter(c["a"], c["b_ind"], s=40, color=ORANGE, edgecolor=PANEL, lw=1, label="independent seeds")
    a.scatter(c["a"], c["b_crn"], s=40, color=BLUE, edgecolor=PANEL, lw=1, label="common random numbers")
    lim = [min(c["a"].min(), c["b_ind"].min(), c["b_crn"].min()) * 0.95, max(c["a"].max(), c["b_ind"].max()) * 1.03]
    a.plot(lim, lim, color=MUTED, lw=1, ls="--")
    a.text(lim[1] - 0.03 * (lim[1] - lim[0]), lim[1] - 0.09 * (lim[1] - lim[0]), "y = x (no effect)",
           color=MUTED, fontsize=8.5, ha="right", rotation=35)
    a.set(xlim=lim, ylim=lim, xlabel="baseline plan: mean Standard wait (min)",
          ylabel="+1 lane 05-07: mean Standard wait (min)", title="Each dot = one simulated day (40 days)")
    a.legend(loc="upper left")

    for y, (key, col, lab) in enumerate((("d_ind", ORANGE, "independent"), ("d_crn", BLUE, "CRN"))):
        m, hw = c[key]
        b.errorbar(m, y, xerr=hw, fmt="o", ms=9, color=col, capsize=7, lw=2.5, mec=PANEL, mew=2)
        b.text(m, y + 0.2, f"{lab}: {m:.3f} +/- {hw:.3f} min", ha="center", color=INK2, fontsize=9.5)
    b.axvline(0, color=MUTED, lw=1)
    b.set_ylim(-0.6, 1.7)
    b.set_yticks([])
    b.set(xlabel="wait saved by the extra lane (min), 95% CI",
          title=f"Same 40 days, {c['var_ratio']:.0f}x smaller variance with CRN")
    save(fig, "crn.png")


def fig_waits(reps):
    fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))
    goal = (30, 10)  # TSA's published wait-time standards
    for k, ax in enumerate(axes):
        w = np.concatenate([r["wait"][r["kind"] == k] for r in reps])
        p95 = np.percentile(w, 95)
        ax.hist(w, bins=np.arange(0, 32, 0.5), density=True, color=STREAM_C[k], alpha=0.85, rwidth=0.85)
        ax.set_yscale("log")
        for x, col, ls, lab in ((p95, INK, "--", f"P95 all day = {p95:.1f} min"),
                                (TARGET_MIN, RED, "-", "our target 15 min (per hour)"),
                                (goal[k], YELLOW, ":", f"TSA standard < {goal[k]} min")):
            ax.axvline(x, color=col, ls=ls, lw=1.6, label=lab)
        ax.set(xlim=(0, 31), xlabel="wait before screening (min)", ylabel="density (log)",
               title=f"{STREAMS[k]}: {len(w):,} passengers, mean {w.mean():.1f} min")
        ax.legend(loc="upper right", fontsize=8.5)
    save(fig, "wait_distributions.png")


# ----------------------------------------------------------------------------------------
def fig_hero(res, plans, reps, summ):
    rob, s = reps["SimOpt robust"], summ["SimOpt robust"]
    fig = plt.figure(figsize=(16, 9))
    gs = fig.add_gridspec(3, 2, width_ratios=[3.2, 1], height_ratios=[1, 2.4, 1],
                          left=0.055, right=0.985, top=0.84, bottom=0.07, hspace=0.1, wspace=0.06)
    fig.text(0.055, 0.935, "ATL CHECKPOINT", size=32, weight="bold")
    fig.text(0.055, 0.885, "Discrete-event simulation of a TSA checkpoint at the world's busiest airport: "
             "NHPP arrivals, ID check, screening lanes, secondary search. 40 replicated days.",
             size=12.5, color=INK2)
    t = np.linspace(T_OPEN, T_END, 1200)
    a0 = fig.add_subplot(gs[0, 0])
    lam = sim.rate(t) * 60
    a0.fill_between(t / 60, 0, lam * (1 - sim.P_PRECHECK), color=BLUE, alpha=0.7, lw=0, label="Standard")
    a0.fill_between(t / 60, lam * (1 - sim.P_PRECHECK), lam, color=AQUA, alpha=0.7, lw=0, label="PreCheck")
    a0.set(ylabel="arrivals / h")
    a0.legend(loc="upper right", ncol=2)
    a0.tick_params(labelbottom=False)
    a1 = fig.add_subplot(gs[1, 0], sharex=a0)
    for k in (0, 1):
        mean_band(a1, np.array([r["nq"][k] for r in rob]), STREAM_C[k], f"{STREAMS[k]} in line", spaghetti=True)
    a1.set(ylabel="passengers in line")
    a1.set_ylim(0, None)
    a1.legend(loc="upper right")
    a1.tick_params(labelbottom=False)
    a2 = fig.add_subplot(gs[2, 0], sharex=a0)
    for k in (0, 1):
        a2.step(*steps(plans["SimOpt robust"][k]), where="post", color=STREAM_C[k], lw=2.2)
    a2.yaxis.set_major_locator(MaxNLocator(integer=True))
    a2.set(ylabel="lanes open")
    a2.set_ylim(0, None)
    hour_axis(a2)

    tiles = fig.add_subplot(gs[:, 1])
    tiles.axis("off")
    m, crn = res["mmc"], res["crn"]
    sipp = summ["SIPP (Erlang C)"]["lane_hours"]
    items = [
        (f"{s['lane_hours']}", "lane-hours / day, sim-optimised", f"SIPP (Erlang C) staffing needs {sipp}"),
        (f"{s['worst0']:.1f} min", "worst-hour P95 wait, Standard",
         f"target 15; {plural(s['viol0'], 'hour')} over; PreCheck {s['worst1']:.1f}"),
        (f"{s['mean0'][0]:.1f} min", "mean wait, Standard", f"PreCheck {s['mean1'][0]:.1f} min"),
        (f"{m['mean']:.3f} min", f"M/M/5 simulated Wq (+/- {m['hw']:.3f})", f"Erlang C exact: {m['wq']:.3f} min"),
        (f"{crn['var_ratio']:.0f}x", "less variance with common", "random numbers (policy comparison)"),
    ]
    for i, (big, lab, sub) in enumerate(items):
        y = 0.99 - i * 0.2
        tiles.add_patch(Rectangle((0.02, y - 0.175), 0.96, 0.17, transform=tiles.transAxes, color=PANEL, lw=0))
        tiles.text(0.08, y - 0.02, big, transform=tiles.transAxes, size=25, weight="bold", va="top")
        tiles.text(0.08, y - 0.095, lab, transform=tiles.transAxes, size=11, color=INK2, va="top")
        tiles.text(0.08, y - 0.13, sub, transform=tiles.transAxes, size=9, color=MUTED, va="top")
    fig.savefig(OUT / "hero.png", dpi=120)
    plt.close(fig)


# ----------------------------------------------------------------------------------------
def make_gif(day, plan, step_min=6, fps=10):
    """Top-down animated checkpoint for one simulated day (sim-optimised staffing)."""
    frames = np.arange(T_OPEN, 22.5 * 60 + 1, step_min)
    counts = {}
    for k in (0, 1):
        m = day["kind"] == k
        ta, tt, st, tl, sl = (day[x][m] for x in ("t_arr", "t_tdc", "s_tdc", "t_lane", "s_lane"))
        # the model has an unlimited buffer between ID check and lanes; draw everyone waiting as one line
        line = lambda g: sim.n_between(ta, tt, g) + sim.n_between(tt + st, tl, g)
        order = np.argsort(tl)
        tl_s, w_s = tl[order], day["wait"][m][order]
        recent = [(tl_s > f - 10) & (tl_s <= f) for f in frames]  # reached a lane in the last 10 min
        counts[k] = dict(line=line(frames), curve=line(sim.MINUTES), tdc=sim.n_between(tt, tt + st, frames),
                         lane=sim.n_between(tl, tl + sl, frames),
                         now_wait=np.array([w_s[r].mean() if r.any() else 0.0 for r in recent]))

    fig = plt.figure(figsize=(10, 5.6))
    ax = fig.add_axes([0.0, 0.0, 0.64, 0.9])
    ax.set_xlim(0, 100); ax.set_ylim(0, 62); ax.axis("off")
    zones = {0: dict(y0=26, rows=14, lane_y=(26, 2.0)), 1: dict(y0=3, rows=6, lane_y=(3.5, 2.3))}
    cols, idle = 26, ("#1f4470", "#155a44")
    dots, pods, lanes, over = {}, {}, {}, {}
    for k, z in zones.items():
        r, c = np.divmod(np.arange(z["rows"] * cols), cols)
        c = np.where(r % 2 == 0, c, cols - 1 - c)                  # serpentine rope line
        z["slots"] = np.c_[45 - c * 1.6, z["y0"] + r * 2.1]
        dots[k] = ax.scatter([], [], s=9, color=STREAM_C[k], lw=0)
        span = z["rows"] * 2.1
        over[k] = ax.text(4.5, z["y0"] + span + 0.2, "", color=INK2, fontsize=8)
        pods[k] = ax.add_collection(PatchCollection(
            [Rectangle((48.5, z["y0"] + (i + 0.5) * span / TDC_PODIUMS[k] - 0.8), 2.2, 1.6) for i in range(TDC_PODIUMS[k])], lw=0))
        y0, dy = z["lane_y"]
        lanes[k] = ax.add_collection(PatchCollection(
            [Rectangle((55, y0 + i * dy), 41, dy * 0.62) for i in range(LANE_MAX[k])], lw=0))
        ax.text(4.5, z["y0"] + span + 2.6, STREAMS[k].upper(), color=STREAM_C[k], fontsize=11, weight="bold")
    for x, lab in ((25, "waiting line"), (49.6, "ID"), (75.5, "screening lanes")):
        ax.text(x, 60.2, lab, color=MUTED, fontsize=8.5, ha="center")
    clock = fig.text(0.02, 0.93, "", fontsize=22, weight="bold")
    fig.text(0.16, 0.945, "One simulated day at an ATL checkpoint (SimPy)", fontsize=11, color=INK2)
    fig.text(0.16, 0.91, "lanes: bright = screening, dim = open & idle, dark = closed", fontsize=8.5, color=MUTED)

    s = fig.add_axes([0.70, 0.42, 0.28, 0.43])
    s.grid(False)
    s.set_xlim(4, 23)
    s.set_ylim(0, max(counts[0]["curve"].max(), counts[1]["curve"].max()) * 1.1)
    s.set_xticks([4, 8, 12, 16, 20], ["04", "08", "12", "16", "20"])
    s.set_title("passengers in line", fontsize=10)
    curves = [s.plot([], [], color=STREAM_C[k], lw=1.6)[0] for k in (0, 1)]
    now = s.axvline(4, color=INK2, lw=0.8)
    fig.text(0.70, 0.30, "          in line  lanes   wait", fontsize=8.5, color=MUTED, family="monospace")
    rows = [fig.text(0.70, 0.24 - 0.06 * k, "", fontsize=8.5, color=STREAM_C[k], family="monospace") for k in (0, 1)]
    fig.text(0.70, 0.08, "wait = mean of passengers who\nreached a lane in the last 10 min", fontsize=7.5, color=MUTED)
    x_min = sim.MINUTES / 60

    def draw(fi):
        f = frames[fi]
        hour = min(int((f - T_OPEN) // 60), H - 1)
        clock.set_text(f"{int(f // 60):02d}:{int(f % 60):02d}")
        for k, z in zones.items():
            c = counts[k]
            n, cap = c["line"][fi], len(z["slots"])
            dots[k].set_offsets(z["slots"][: min(n, cap)] if n else np.empty((0, 2)))
            over[k].set_text(f"+{n - cap} more" if n > cap else "")
            pods[k].set_facecolor([STREAM_C[k] if i < c["tdc"][fi] else GRID for i in range(TDC_PODIUMS[k])])
            n_open, n_busy = plan[k, hour], c["lane"][fi]
            lanes[k].set_facecolor([STREAM_C[k] if i < n_busy else idle[k] if i < n_open else "#1e1e1c"
                                    for i in range(LANE_MAX[k])])
            m = x_min <= f / 60
            curves[k].set_data(x_min[m], c["curve"][m])
            rows[k].set_text(f"{STREAMS[k]:9s}{n:6d}{n_open:7d}{c['now_wait'][fi]:6.1f} min")
        now.set_xdata([f / 60, f / 60])
        return []

    anim = FuncAnimation(fig, draw, frames=len(frames), blit=False)
    anim.save(OUT / "checkpoint.gif", writer=PillowWriter(fps=fps), dpi=72, savefig_kwargs={"facecolor": BG})
    plt.close(fig)


def make_all(res, plans, reps, summ, traces, gif_day):
    fig_nhpp(res["nhpp"])
    fig_erlang(res["mmc"])
    fig_warmup(res["warmup"], res["mmc"]["wq"])
    fig_staffing(plans, summ, traces)
    fig_policies(summ, reps)
    fig_queue(reps)
    fig_crn(res["crn"])
    fig_waits(reps["SimOpt robust"])
    fig_hero(res, plans, reps, summ)
    make_gif(gif_day, plans["SimOpt robust"])

"""L03 ATL Checkpoint: discrete-event simulation of a TSA checkpoint at Hartsfield-Jackson.

Run from the repo root:  uv run python L03_atl_checkpoint/sim.py
"""
import os

os.environ.setdefault("OMP_NUM_THREADS", "4")

import time
from collections import deque
from heapq import heappop, heappush
from multiprocessing import Pool
from pathlib import Path

import numpy as np
import simpy
from scipy import stats

OUT = Path(__file__).parent / "output"
WORKERS = 4  # CPU etiquette: shared machine

# ----------------------------------------------------------------------------------------
# Model parameters. (P) = anchored to a public figure, (A) = assumption. See README.
# Time unit: minutes after midnight.
# ----------------------------------------------------------------------------------------
T_OPEN, H = 4 * 60, 20            # lanes are staffed in 20 hourly blocks, 04:00-24:00 (A)
T_END = T_OPEN + 60 * H
T_LAST = 22 * 60                  # last passenger arrives before 22:00 (A)
DAILY_PAX = 22_000                # passengers/day at ONE checkpoint (A; ATL screens ~82k/day over several checkpoints, P)
P_PRECHECK = 0.30                 # share of passengers in the PreCheck stream (A)
LANE_MEAN = (60 / 150, 60 / 300)  # min/passenger per lane: ~150/h standard, ~300/h PreCheck (P)
LANE_SHAPE = 2.0                  # gamma shape of per-passenger lane time -> CV 0.71 (A)
TDC_MEAN, TDC_SHAPE = 10 / 60, 3.0  # travel-document check: 10 s mean, CV 0.58 (A)
TDC_PODIUMS = (7, 3)              # document-check podiums (standard, PreCheck), open all day (A)
LANE_MAX = (16, 6)                # physical lanes (standard, PreCheck) (A)
Q_SECONDARY = (0.10, 0.05)        # P(secondary bag search) (standard, PreCheck) (A)
SEC_MEAN, SEC_SHAPE, SEC_OFFICERS = 3.0, 2.0, 12  # secondary search: 3 min mean, 12 officers (A)
TARGET_MIN = 15.0                 # service level: P95 wait <= 15 min for arrivals in EVERY hour
STREAMS = ("Standard", "PreCheck")

# Arrival-rate shape: Delta's departure banks, seen ~1-2 h earlier at the checkpoint (A).
BUMPS = [(5.8, 1.00, 0.55), (8.4, 0.50, 0.8), (10.9, 0.62, 0.8),
         (13.8, 0.64, 0.9), (16.6, 0.70, 0.8), (18.9, 0.40, 0.8)]  # (hour, height, width h)


def _shape(t):
    h = np.asarray(t) / 60
    s = 0.10 + sum(a * np.exp(-0.5 * ((h - m) / w) ** 2) for m, a, w in BUMPS)
    return s / (1 + np.exp(-(h - 4.3) * 10)) / (1 + np.exp((h - 21.4) * 5))


_GRID = np.linspace(T_OPEN, T_LAST, 20_001)
_SCALE = DAILY_PAX / np.trapezoid(_shape(_GRID), _GRID)


def rate(t):
    """Total arrival rate lambda(t) at the checkpoint, passengers/min."""
    t = np.asarray(t, float)
    return _SCALE * _shape(t) * ((t >= T_OPEN) & (t < T_LAST))


LAM_MAX = 1.01 * rate(_GRID).max()  # thinning needs lambda(t) <= LAM_MAX everywhere


def expected_per_hour():
    """Lambda(hour block) = integral of lambda(t) over each staffing hour, shape (H,)."""
    return np.array([np.trapezoid(rate(g), g) for g in
                     (np.linspace(T_OPEN + 60 * h, T_OPEN + 60 * (h + 1), 601) for h in range(H))])


# ----------------------------------------------------------------------------------------
# Random variate generation
# ----------------------------------------------------------------------------------------
def nhpp_thinning(rng, rate_fn, t0, t1, lam_max):
    """Lewis-Shedler thinning: candidate points from a homogeneous Poisson process at lam_max,
    each kept with probability rate(t)/lam_max. (Given the count, HPP points are iid uniform.)"""
    n = rng.poisson(lam_max * (t1 - t0))
    cand = np.sort(rng.uniform(t0, t1, n))
    y = rng.uniform(size=n) * lam_max          # accept iff the point (t, y) lies under rate(t)
    return cand, y, y < rate_fn(cand)


def make_day(seed):
    """Pre-generate every random input of one day. All staffing plans then see the SAME
    passengers with the SAME service needs: perfectly synchronised common random numbers."""
    rng = np.random.default_rng(seed)
    cand, _, keep = nhpp_thinning(rng, rate, T_OPEN, T_LAST, LAM_MAX)
    t = cand[keep]
    n = len(t)
    kind = (rng.uniform(size=n) < P_PRECHECK).astype(np.int8)  # 0 standard, 1 PreCheck
    lane_mean = np.where(kind == 1, LANE_MEAN[1], LANE_MEAN[0])
    return dict(
        t_arr=t, kind=kind,
        s_tdc=rng.gamma(TDC_SHAPE, TDC_MEAN / TDC_SHAPE, n),
        s_lane=rng.gamma(LANE_SHAPE, lane_mean / LANE_SHAPE),
        sec=rng.uniform(size=n) < np.where(kind == 1, Q_SECONDARY[1], Q_SECONDARY[0]),
        s_sec=rng.gamma(SEC_SHAPE, SEC_MEAN / SEC_SHAPE, n),
    )


# ----------------------------------------------------------------------------------------
# The checkpoint in SimPy (process-interaction worldview)
# ----------------------------------------------------------------------------------------
def lane_controller(env, lanes, plan):
    """Keep plan[h] lanes open during hour h. A closing lane finishes its current passenger first
    (a Resource only grants new requests while users < capacity)."""
    for h, n_open in enumerate(plan):
        yield env.timeout(T_OPEN + 60 * h - env.now)
        # ponytail: SimPy has no public capacity setter; the documented alternative (PriorityResource
        # + high-priority 'closed' tokens) gives identical results but re-sorts the queue per request.
        lanes._capacity = int(n_open)
        for _ in range(max(0, lanes._capacity - len(lanes.users))):
            lanes._trigger_put(None)                 # newly opened lanes pull waiting passengers now


def run_day(day, plan, t_end=T_END):
    """Simulate one day under a staffing plan (array 2 x H of open lanes)."""
    env = simpy.Environment()
    tdc = [simpy.Resource(env, c) for c in TDC_PODIUMS]
    lanes = [simpy.Resource(env, 1) for _ in (0, 1)]  # capacity set hourly by lane_controller
    officers = simpy.Resource(env, SEC_OFFICERS)
    # plain Python floats: numpy scalars would make every SimPy clock operation ~2x slower
    t_arr, kind, s_tdc, s_lane, sec, s_sec = (day[k].tolist() for k in ("t_arr", "kind", "s_tdc", "s_lane", "sec", "s_sec"))
    n = len(t_arr)
    t_tdc, t_lane, t_done = np.full(n, np.nan), np.full(n, np.nan), np.full(n, np.nan)

    def passenger(i, k):
        with tdc[k].request() as r:          # 1. travel-document check
            yield r
            t_tdc[i] = env.now
            yield env.timeout(s_tdc[i])
        with lanes[k].request() as r:        # 2. screening lane (X-ray + body scanner)
            yield r
            t_lane[i] = env.now
            yield env.timeout(s_lane[i])
        if sec[i]:                           # 3. secondary bag search with probability q
            with officers.request() as r:
                yield r
                yield env.timeout(s_sec[i])
        t_done[i] = env.now

    def source():
        for i in range(n):
            yield env.timeout(max(0.0, t_arr[i] - env.now))
            env.process(passenger(i, kind[i]))

    for k in (0, 1):
        env.process(lane_controller(env, lanes[k], plan[k]))
    env.process(source())
    env.run(until=t_end)
    # wait = time in queues before screening starts; passengers still queued at t_end are censored
    wait = np.where(np.isnan(t_lane), t_end - day["t_arr"], t_lane - day["t_arr"] - day["s_tdc"])
    return dict(day, t_tdc=t_tdc, t_lane=t_lane, t_done=t_done, wait=wait)


MINUTES = np.arange(T_OPEN, T_END + 1, 1.0)


def n_between(a, b, grid):
    """How many intervals [a_i, b_i) cover each grid time (NaN end = still open)."""
    b = np.where(np.isnan(b), np.inf, b)
    return np.searchsorted(np.sort(a), grid, "right") - np.searchsorted(np.sort(b), grid, "right")


def _day_task(args):
    """Pool worker: one replication, returned as a compact summary."""
    seed, plan, t_end = args
    r = run_day(make_day(seed), plan, t_end)
    out = {k: r[k] for k in ("t_arr", "kind", "wait")}
    out["sys"] = r["t_done"] - r["t_arr"]
    out["nq"] = np.array([n_between(r["t_arr"][m], r["t_tdc"][m], MINUTES)
                          + n_between((r["t_tdc"] + r["s_tdc"])[m], r["t_lane"][m], MINUTES)
                          for m in (r["kind"] == 0, r["kind"] == 1)])  # passengers waiting, per stream
    return out


def simulate(plan, seeds, t_end=T_END):
    return POOL.map(_day_task, [(s, plan, t_end) for s in seeds])


# ----------------------------------------------------------------------------------------
# Queueing theory (the yardstick)
# ----------------------------------------------------------------------------------------
def erlang_c(c, a):
    """P(an arrival waits) in M/M/c with offered load a = lam/mu, via the stable Erlang-B recursion."""
    b = 1.0
    for k in range(1, c + 1):
        b = a * b / (k + a * b)
    return c * b / (c - a * (1 - b))


def mmc_wq(lam, mu, c):
    return erlang_c(c, lam / mu) / (c * mu - lam)


def mmc_p_wait_gt(t, lam, mu, c):
    """P(Wq > t) = C(c, a) * exp(-(c mu - lam) t)."""
    return erlang_c(c, lam / mu) * np.exp(-(c * mu - lam) * np.asarray(t))


# ----------------------------------------------------------------------------------------
# M/M/c in both DES worldviews. Same inputs -> identical waits.
# ----------------------------------------------------------------------------------------
MMC_LAM, MMC_MU, MMC_C = 11.0, 2.5, 5  # 5 lanes x 150/h, 660 passengers/h -> rho = 0.88


def mmc_inputs(seed, n, lam, mu):
    rng = np.random.default_rng(seed)
    return np.cumsum(rng.exponential(1 / lam, n)), rng.exponential(1 / mu, n)


def mmc_simpy(arr, svc, c):
    """Process interaction: write the life of ONE customer; SimPy owns the event list."""
    arr, svc = arr.tolist(), svc.tolist()
    env = simpy.Environment()
    servers = simpy.Resource(env, c)
    wait = np.empty(len(arr))

    def customer(i):
        with servers.request() as req:
            yield req
            wait[i] = env.now - arr[i]
            yield env.timeout(svc[i])

    def source():
        for i, t in enumerate(arr):
            yield env.timeout(t - env.now)
            env.process(customer(i))

    env.process(source())
    env.run()
    return wait


def mmc_events(arr, svc, c):
    """Event scheduling: a future event list (heap) and one routine per event type."""
    arr, svc = arr.tolist(), svc.tolist()
    ARRIVE, DEPART = 0, 1
    fel, queue, busy, wait = [(arr[0], ARRIVE, 0)], deque(), 0, np.empty(len(arr))
    while fel:
        t, ev, i = heappop(fel)
        if ev == ARRIVE:
            if i + 1 < len(arr):
                heappush(fel, (arr[i + 1], ARRIVE, i + 1))   # bootstrap the next arrival
            if busy < c:
                busy += 1
                wait[i] = 0.0
                heappush(fel, (t + svc[i], DEPART, i))
            else:
                queue.append(i)
        elif queue:                                          # departure: next in line starts
            j = queue.popleft()
            wait[j] = t - arr[j]
            heappush(fel, (t + svc[j], DEPART, j))
        else:
            busy -= 1
    return wait


def _mmc_task(args):
    seed, n, lam, engine = args
    arr, svc = mmc_inputs(seed, n, lam, MMC_MU)
    return (mmc_simpy if engine == "simpy" else mmc_events)(arr, svc, MMC_C)


# ----------------------------------------------------------------------------------------
# Output analysis
# ----------------------------------------------------------------------------------------
def t_ci(x, axis=0):
    """Mean and 95% t-based half-width across independent replications."""
    x = np.asarray(x, float)
    n = x.shape[axis]
    return x.mean(axis), stats.t.ppf(0.975, n - 1) * x.std(axis, ddof=1) / np.sqrt(n)


def welch_ma(y, w):
    """Welch's moving average: centred window of 2w+1, shrunk to 2i+1 near the start."""
    c = np.concatenate([[0.0], np.cumsum(y)])
    i = np.arange(len(y) - w)
    s = np.minimum(i, w)
    return (c[i + s + 1] - c[i - s]) / (2 * s + 1)


COVER_REPS, COVER_NS = 10, (800, 1600, 3200, 6400)  # one CI = 10 replications from empty


def _coverage_task(args):
    """One macro-replication: for each run length n, does the 95% CI cover the true Wq
    without deletion / after deleting the first d customers? -> bool array (len(NS), 2)."""
    seed, d = args
    wq = mmc_wq(MMC_LAM, MMC_MU, MMC_C)
    ws = [_mmc_task((seed * 1000 + r, max(COVER_NS), MMC_LAM, "events")) for r in range(COVER_REPS)]
    out = []
    for n in COVER_NS:  # a run of n customers = the first n customers of the long run
        for dd in (0, d):
            m, hw = t_ci([w[dd:n].mean() for w in ws])
            out.append(abs(m - wq) <= hw)
    return np.reshape(out, (len(COVER_NS), 2))


def exceed(reps):
    """Per-replication fraction of passengers arriving in each hour who wait > 15 min
    -> (R, 2, H), NaN where nobody arrives. 'P95 <= 15 min' <=> 'fraction > 15 min <= 5%'."""
    out = np.full((len(reps), 2, H), np.nan)
    for i, r in enumerate(reps):
        hour = ((r["t_arr"] - T_OPEN) // 60).astype(int)
        for k in (0, 1):
            m = r["kind"] == k
            n = np.bincount(hour[m], minlength=H)[:H]
            bad = np.bincount(hour[m], weights=r["wait"][m] > TARGET_MIN, minlength=H)[:H]
            out[i, k] = np.where(n > 0, bad / np.maximum(n, 1), np.nan)
    return out


def hourly_p95(reps):
    """P95 wait (min) of passengers ARRIVING in each hour, per stream -> (2, H).
    Pools the replications: the sample-average approximation the optimiser uses."""
    t, k, w = (np.concatenate([r[x] for r in reps]) for x in ("t_arr", "kind", "wait"))
    hour = ((t - T_OPEN) // 60).astype(int)
    out = np.full((2, H), np.nan)
    for s in (0, 1):
        for h in np.unique(hour[k == s]):
            out[s, h] = np.percentile(w[(k == s) & (hour == h)], 95)
    return out


def summarize(reps, plan):
    """Headline numbers of one policy over its replications."""
    pooled = hourly_p95(reps)
    ex_m, ex_hw = t_ci(exceed(reps))
    s = dict(lane_hours=int(plan.sum()), hourly=pooled, exceed=(ex_m, ex_hw))
    for k in (0, 1):
        h = np.nanargmax(ex_m[k])
        s[f"exc{k}"] = (ex_m[k, h], ex_hw[k, h], h)   # worst hour: P(wait > 15 min) and its CI
        s[f"sys{k}"] = t_ci([np.nanmean(r["sys"][r["kind"] == k]) for r in reps])
        s[f"mean{k}"] = t_ci([r["wait"][r["kind"] == k].mean() for r in reps])
        s[f"p95_{k}"] = t_ci([np.percentile(r["wait"][r["kind"] == k], 95) for r in reps])
        s[f"worst{k}"] = np.nanmax(pooled[k])
        s[f"viol{k}"] = int(np.nansum(pooled[k] > TARGET_MIN))
    return s


# ----------------------------------------------------------------------------------------
# Staffing: analytic SIPP start, then simulation optimisation
# ----------------------------------------------------------------------------------------
def sipp_plan():
    """Stationary Independent Period-by-Period: treat each hour as a steady-state M/M/c at that
    hour's average rate; pick the fewest lanes with P(Wq > 15 min) <= 5%."""
    lam_h = expected_per_hour()
    plan = np.ones((2, H), int)
    for k, share in ((0, 1 - P_PRECHECK), (1, P_PRECHECK)):
        mu = 1 / LANE_MEAN[k]
        for h in range(H):
            lam = lam_h[h] * share / 60
            c = max(1, int(lam / mu) + 1)
            while c < LANE_MAX[k] and mmc_p_wait_gt(TARGET_MIN, lam, mu, c) > 0.05:
                c += 1
            plan[k, h] = c
    return plan


def optimise(start, seeds, z):
    """Greedy forward sweep over hours (waits in hour h depend only on lanes up to ~h+1).
    For each hour and stream, walk the lane count down (or up) until the constraint flips.
    Constraint: mean + z * (95% half-width) of the per-rep exceedance fraction <= 5%
    (z=0: naive sample average; z=1: statistically conservative upper bound).
    Every candidate is scored on the SAME seeds (CRN), so comparisons are paired.
    ponytail: greedy per-hour search, not a proven optimum; ranking-and-selection or an
    integer program on a fitted delay model would certify it."""
    plan = start.copy()
    trace = []  # (evaluations so far, lane-hours of current plan)
    for h in range(H):
        t_end = T_OPEN + 60 * min(H, h + 2)
        best, step = [None, None], [0, 0]
        while None in best:
            m, hw = t_ci(exceed(simulate(plan, seeds, t_end))[:, :, : h + 1])
            bound = np.nan_to_num(m + z * hw)
            for k in (0, 1):
                if best[k] is not None:
                    continue
                n, ok = plan[k, h], bool(np.all(bound[k] <= 0.05))
                if ok and (step[k] == 1 or n == 1):
                    best[k] = n
                elif ok:
                    step[k], plan[k, h] = -1, n - 1
                elif step[k] == -1 or n == LANE_MAX[k]:
                    best[k] = plan[k, h] = min(n + 1, LANE_MAX[k])
                else:
                    step[k], plan[k, h] = 1, n + 1
            trace.append((len(trace) + 1, int(plan.sum())))
    return plan, np.array(trace)


# ----------------------------------------------------------------------------------------
# Experiments
# ----------------------------------------------------------------------------------------
def exp_nhpp(n_days=300):
    rng = np.random.default_rng(3)
    bins = np.arange(T_OPEN, T_LAST + 1, 15.0)
    counts, totals = [], []
    for _ in range(n_days):
        cand, _, keep = nhpp_thinning(rng, rate, T_OPEN, T_LAST, LAM_MAX)
        counts.append(np.histogram(cand[keep], bins)[0])
        totals.append(keep.sum())
    expected = np.array([np.trapezoid(rate(g), g) for g in (np.linspace(a, a + 15, 101) for a in bins[:-1])])
    cand, y, keep = nhpp_thinning(np.random.default_rng(4), rate, T_OPEN, T_LAST, LAM_MAX)
    return dict(bins=bins, counts=np.array(counts), expected=expected, totals=np.array(totals),
                Lambda=np.trapezoid(rate(_GRID), _GRID), cand=cand, y=y, keep=keep)


def exp_mmc():
    """Erlang-C verification: SimPy vs event scheduling, and a rho sweep."""
    n, d, R = 22_000, 2_000, 20
    arr, svc = mmc_inputs(0, n, MMC_LAM, MMC_MU)
    t0 = time.perf_counter(); w_sp = mmc_simpy(arr, svc, MMC_C); t_sp = time.perf_counter() - t0
    t0 = time.perf_counter(); w_ev = mmc_events(arr, svc, MMC_C); t_ev = time.perf_counter() - t0
    reps = POOL.map(_mmc_task, [(s, n, MMC_LAM, "simpy") for s in range(R)])
    m, hw = t_ci([w[d:].mean() for w in reps])
    sweep = []
    for rho in (0.5, 0.65, 0.8, 0.88, 0.94):
        lam = rho * MMC_C * MMC_MU
        ws = POOL.map(_mmc_task, [(100 + s, n, lam, "events") for s in range(R)])
        sweep.append((rho, *t_ci([w[d:].mean() for w in ws]), mmc_wq(lam, MMC_MU, MMC_C)))
    surv_t = np.linspace(0, 7, 141)
    surv = np.array([1 - np.searchsorted(np.sort(w[d:]), surv_t, "right") / (n - d) for w in reps])
    return dict(wq=mmc_wq(MMC_LAM, MMC_MU, MMC_C), mean=m, hw=hw, surv_t=surv_t, surv=surv,
                pw=t_ci([(w[d:] > 0).mean() for w in reps]), pc=erlang_c(MMC_C, MMC_LAM / MMC_MU),
                engines_maxdiff=float(np.max(np.abs(w_sp - w_ev))), t_simpy=t_sp, t_events=t_ev,
                sweep=np.array(sweep), R=R, n=n, d=d)


def exp_warmup(R=500, n=3_000, w=50, macro=300):
    """Welch's procedure on M/M/c started empty, then the payoff: CI coverage with/without deletion."""
    ybar = np.mean(POOL.map(_mmc_task, [(10_000 + s, n, MMC_LAM, "events") for s in range(R)]), 0)
    ma = welch_ma(ybar, w)
    plateau = ma[len(ma) // 2:].mean()
    d = int(np.argmax(ma >= 0.95 * plateau))
    d = int(np.ceil(1.5 * d / 50) * 50)  # Welch is an eyeball method: round up with a safety factor
    cov = np.array(POOL.map(_coverage_task, [(20_000 + s, d) for s in range(macro)]))
    return dict(ybar=ybar, ma=ma, w=w, d=d, plateau=plateau, coverage=cov.mean(0), macro=macro)


def exp_crn(opt, base_reps, val_seeds, ind_seeds):
    """Is one extra standard lane in the 05:00-07:00 bank worth it? CRN vs independent sampling."""
    alt = opt.copy()
    alt[0, 1:3] = np.minimum(alt[0, 1:3] + 1, LANE_MAX[0])
    mean_std = lambda reps: np.array([r["wait"][r["kind"] == 0].mean() for r in reps])
    a = mean_std(base_reps)
    b_crn, b_ind = mean_std(simulate(alt, val_seeds)), mean_std(simulate(alt, ind_seeds))
    return dict(a=a, b_crn=b_crn, b_ind=b_ind, d_crn=t_ci(a - b_crn), d_ind=t_ci(a - b_ind),
                var_ratio=np.var(a - b_ind, ddof=1) / np.var(a - b_crn, ddof=1),
                corr=np.corrcoef(a, b_crn)[0, 1])


# ----------------------------------------------------------------------------------------
def verify(res):
    """Assertions against known results; printed numbers are the real ones from this run."""
    m = res["mmc"]
    print("\n=== VERIFICATION ===")
    print(f"[1] SimPy vs event scheduling, same {m['n']} inputs: max |dW| = {m['engines_maxdiff']:.1e} min "
          f"(SimPy {m['t_simpy']:.2f} s, heapq {m['t_events']:.3f} s)")
    assert m["engines_maxdiff"] < 1e-9
    print(f"[2] M/M/{MMC_C} rho={MMC_LAM / MMC_C / MMC_MU:.2f}: Erlang-C Wq = {m['wq']:.4f} min, "
          f"SimPy {m['R']} reps = {m['mean']:.4f} +/- {m['hw']:.4f}")
    assert abs(m["mean"] - m["wq"]) <= m["hw"], "Erlang-C Wq outside the 95% CI"
    print(f"    P(wait>0): Erlang-C {m['pc']:.4f}, sim {m['pw'][0]:.4f} +/- {m['pw'][1]:.4f}")
    for rho, mm, hw, wq in m["sweep"]:
        print(f"    rho={rho:.2f}: Wq theory {wq:.4f}, sim {mm:.4f} +/- {hw:.4f}  {'ok' if abs(mm - wq) <= hw else 'MISS'}")
    n = res["nhpp"]
    z = (n["totals"].mean() - n["Lambda"]) / (n["totals"].std(ddof=1) / np.sqrt(len(n["totals"])))
    print(f"[3] NHPP thinning, {len(n['totals'])} days: daily count mean {n['totals'].mean():.1f} vs Lambda "
          f"{n['Lambda']:.1f} (z={z:+.2f}), var/mean = {n['totals'].var(ddof=1) / n['totals'].mean():.3f} (Poisson: 1)")
    assert abs(z) < 3.5
    wu = res["warmup"]
    print(f"[4] Welch warm-up d = {wu['d']} customers. 95% CI coverage over {wu['macro']} macro-reps (nominal 0.95):")
    for n_, (c0, c1) in zip(COVER_NS, wu["coverage"]):
        print(f"    run length {n_:5d}: no deletion {c0:.3f}, delete first {wu['d']}: {c1:.3f}")
    assert np.all(wu["coverage"][:, 1] >= wu["coverage"][:, 0] - 0.01)


POOL = None


def main():
    global POOL
    OUT.mkdir(exist_ok=True)
    T0 = time.perf_counter()
    tick = lambda s: print(f"  [{time.perf_counter() - T0:6.1f}s] {s}", flush=True)
    res = {}
    with Pool(WORKERS) as POOL:
        res["nhpp"] = exp_nhpp(); tick("NHPP thinning")
        res["mmc"] = exp_mmc(); tick("M/M/c vs Erlang C")
        res["warmup"] = exp_warmup(); tick("Welch warm-up + CI coverage")
        verify(res)

        sipp = sipp_plan()
        val_seeds, ind_seeds = range(5000, 5040), range(7000, 7040)
        naive, tr_naive = optimise(sipp, range(1000, 1012), z=0)
        tick(f"sim-optimisation, naive: 12 reps, point estimate ({len(tr_naive)} evaluations)")
        robust, tr_robust = optimise(sipp, range(1000, 1024), z=1)
        tick(f"sim-optimisation, robust: 24 reps, upper confidence bound ({len(tr_robust)} evaluations)")
        flat = np.array([[int(np.ceil(robust[k].sum() / H))] * H for k in (0, 1)])
        plans = {"SIPP (Erlang C)": sipp, "SimOpt naive": naive, "SimOpt robust": robust, "Flat": flat}
        reps = {name: simulate(p, val_seeds) for name, p in plans.items()}
        tick("validation, 40 fresh reps x 4 policies")
        summ = {name: summarize(reps[name], plans[name]) for name in plans}
        res["crn"] = exp_crn(robust, reps["SimOpt robust"], val_seeds, ind_seeds); tick("CRN experiment")

    print("\n=== STAFFING POLICIES (validated on 40 fresh seeds; waits in min) ===")
    print(f"{'policy':16s} lane-h   mean wait std / pre      day P95 std / pre   worst-hour P95 std / pre  hours>15"
          "   worst-hour P(wait>15) std")
    for name, s in summ.items():
        e = s["exc0"]
        print(f"{name:16s} {s['lane_hours']:5d}   {s['mean0'][0]:6.2f}+/-{s['mean0'][1]:.2f} / {s['mean1'][0]:5.2f}+/-{s['mean1'][1]:.2f}"
              f"   {s['p95_0'][0]:6.2f} / {s['p95_1'][0]:5.2f}      {s['worst0']:6.2f} / {s['worst1']:6.2f}         {s['viol0']}/{s['viol1']}"
              f"      {100 * e[0]:5.1f}% +/- {100 * e[1]:.1f}% ({4 + e[2]:02d}:00)")
    s = summ["SimOpt robust"]
    print(f"SimOpt robust, time in system incl. secondary search: std {s['sys0'][0]:.2f} +/- {s['sys0'][1]:.2f} min, "
          f"pre {s['sys1'][0]:.2f} +/- {s['sys1'][1]:.2f} min")
    for name, p in plans.items():
        print(f"{name:16s} std {p[0].tolist()}\n{'':16s} pre {p[1].tolist()}")
    c = res["crn"]
    print(f"\nCRN: +1 standard lane 05-07 cuts mean standard wait by {c['d_crn'][0]:.3f} +/- {c['d_crn'][1]:.3f} min (CRN) "
          f"vs {c['d_ind'][0]:.3f} +/- {c['d_ind'][1]:.3f} (independent); variance ratio {c['var_ratio']:.0f}x, corr {c['corr']:.4f}")

    import plots  # figures live in their own file
    gif_day = run_day(make_day(val_seeds[0]), robust)
    plots.make_all(res, plans, reps, summ, dict(naive=tr_naive, robust=tr_robust), gif_day)
    tick("figures + GIF")
    print(f"\nTotal runtime {time.perf_counter() - T0:.1f} s")


if __name__ == "__main__":
    main()