#!/usr/bin/env python3
"""
Reproducible simulations for "The Orchestration Substrate" (blog.pdj.dev).

Generates two figures (themed SVG, transparent background for the dark site):
  1. orchestration-scheduling.svg - mean weighted completion time per policy,
                                     showing the c-mu rule is optimal.
  2. orchestration-durability.svg - reliability decay (exponential) vs durable
                                     expected work (linear), closed form vs Monte Carlo.

Run:  python3 simulations.py
Deterministic (seeded). Requires numpy, matplotlib. No network, no data files.
"""

from __future__ import annotations
import pathlib
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

SEED = 7
RNG = np.random.default_rng(SEED)
OUT = pathlib.Path(__file__).resolve().parents[2] / "assets" / "figures"
OUT.mkdir(parents=True, exist_ok=True)

# ---- dark-site theme -------------------------------------------------------
ACCENT, C2, C3, C4 = "#4da3ff", "#ff7ab2", "#82d9ff", "#d2a8ff"
TEXT, MUTED, GRID = "#c9c9cf", "#8a8a90", "#2c2c2e"
plt.rcParams.update({
    "figure.facecolor": "none", "axes.facecolor": "none", "savefig.facecolor": "none",
    "font.family": "DejaVu Sans", "font.size": 11,
    "text.color": TEXT, "axes.labelcolor": TEXT,
    "xtick.color": MUTED, "ytick.color": MUTED,
    "axes.edgecolor": GRID, "grid.color": GRID,
    "axes.spines.top": False, "axes.spines.right": False,
    "legend.frameon": False, "figure.dpi": 110,
})

def style(ax):
    ax.grid(True, alpha=0.5, linewidth=0.7)
    ax.tick_params(length=3)
    for s in ("left", "bottom"):
        ax.spines[s].set_linewidth(0.8)

def save(fig, name):
    fig.tight_layout()
    fig.savefig(OUT / name, format="svg", transparent=True, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {OUT / name}")

# ---------------------------------------------------------------------------
# Figure 1: the c-mu rule minimizes mean weighted completion time.
# On a single server, process N tasks with holding costs c_i and service
# times S_i (so mu_i = 1/E[S_i]). For each random instance, compute the total
# weighted completion time  sum_i c_i * C_i  under four orderings:
#   - random order
#   - shortest processing time first  (largest mu_i, i.e. smallest S_i)
#   - highest holding cost first      (largest c_i)
#   - c-mu                            (largest c_i / S_i == c_i * mu_i)
# Average over many instances. c-mu is provably lowest.
# ---------------------------------------------------------------------------
def weighted_completion(order, c, s):
    """Total weighted completion time when tasks are served in `order`."""
    t = 0.0
    total = 0.0
    for i in order:
        t += s[i]            # completion time of task i on a single server
        total += c[i] * t    # weighted by its holding-cost rate
    return total

def fig_scheduling():
    N = 8                    # tasks per instance
    instances = 4000         # Monte Carlo over random instances
    labels = ["random", "shortest first", "highest cost first", "c-mu rule"]
    colors = [MUTED, C3, C4, ACCENT]
    totals = {k: np.empty(instances) for k in labels}

    for m in range(instances):
        c = RNG.uniform(0.2, 5.0, N)          # holding-cost rates c_i
        s = RNG.uniform(0.2, 5.0, N)          # service times S_i (mu_i = 1/s)
        idx = np.arange(N)

        rand_order = RNG.permutation(N)
        spt_order = idx[np.argsort(s)]                 # smallest service time first
        hc_order = idx[np.argsort(-c)]                 # largest cost first
        cmu_order = idx[np.argsort(-(c / s))]          # largest c_i * mu_i first

        totals["random"][m] = weighted_completion(rand_order, c, s)
        totals["shortest first"][m] = weighted_completion(spt_order, c, s)
        totals["highest cost first"][m] = weighted_completion(hc_order, c, s)
        totals["c-mu rule"][m] = weighted_completion(cmu_order, c, s)

    means = [totals[k].mean() for k in labels]
    sems = [totals[k].std(ddof=1) / np.sqrt(instances) for k in labels]

    fig, ax = plt.subplots(figsize=(7.0, 3.9))
    x = np.arange(len(labels))
    ax.bar(x, means, yerr=sems, color=colors, width=0.62,
           error_kw=dict(ecolor=TEXT, elinewidth=1.0, capsize=4), zorder=2)
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.set_ylabel("mean weighted completion time")
    ax.set_ylim(0, max(means) * 1.12)
    style(ax)
    for xi, mv in zip(x, means):
        ax.text(float(xi), mv + max(means) * 0.02, f"{mv:.1f}",
                ha="center", va="bottom", color=TEXT, fontsize=9.5)
    save(fig, "orchestration-scheduling.svg")
    return means

# ---------------------------------------------------------------------------
# Figure 2: durability linearizes cost.
# Naive run: a horizon-n pipeline aborts on any step failure, so end-to-end
# success is p_bar^n (exponential decay). Restart-on-failure repeats the WHOLE
# run until it completes, so the expected number of full-run attempts is
# 1 / p_bar^n, and expected total step-work is n / p_bar^n (exponential in n).
# Durable run: checkpoint each step; attempts at step i are geometric with
# mean 1/p_i, so expected total work is sum_i 1/p_i = n / p_bar (linear in n).
# Overlay Monte Carlo on both closed forms; log scale on the work axis.
# ---------------------------------------------------------------------------
def fig_durability():
    pbar = 0.99
    ns = np.arange(1, 401)

    # --- panel A: end-to-end success collapses exponentially ---
    succ = pbar ** ns

    # --- panel B: expected total work, naive restart vs durable retry ---
    work_naive = ns / (pbar ** ns)     # restart the whole run on any failure
    work_durable = ns / pbar           # checkpoint and retry the failed step

    fig, (axL, axR) = plt.subplots(1, 2, figsize=(7.4, 3.9))

    # left: reliability decay with Monte Carlo
    axL.plot(ns, succ, color=C2, lw=1.8, label="closed form  p̄ⁿ")
    sub = ns[::40]
    trials = 4000
    mc_succ = []
    for n in sub:
        steps = RNG.random((trials, n)) < pbar     # each step succeeds?
        mc_succ.append(steps.all(axis=1).mean())   # whole run succeeds only if all do
    axL.scatter(sub, mc_succ, color=C2, s=18, zorder=3, edgecolor="none")
    axL.set_xlabel("horizon, n (steps)")
    axL.set_ylabel("naive end-to-end P(success)")
    axL.set_ylim(0, 1.02)
    axL.set_xlim(1, 400)
    axL.legend(loc="best")
    style(axL)

    # right: expected total work, log scale to show exponential vs linear
    axR.plot(ns, work_naive, color=C4, lw=1.8, label="naive restart-on-failure")
    axR.plot(ns, work_durable, color=ACCENT, lw=1.8, label="durable checkpoint-retry")
    sub2 = ns[::40]
    mc_naive, mc_durable = [], []
    rep = 600
    for n in sub2:
        # naive: count full-run attempts until one run has all n steps succeed
        attempts = []
        for _ in range(rep):
            k = 0
            while True:
                k += 1
                if (RNG.random(n) < pbar).all():
                    break
            attempts.append(k * n)         # each attempt does n steps of work
        mc_naive.append(np.mean(attempts))
        # durable: per step, geometric retries until that step succeeds
        work = []
        for _ in range(rep):
            w = 0
            for _ in range(n):
                while True:
                    w += 1
                    if RNG.random() < pbar:
                        break
            work.append(w)
        mc_durable.append(np.mean(work))
    axR.scatter(sub2, mc_naive, color=C4, s=18, zorder=3, edgecolor="none")
    axR.scatter(sub2, mc_durable, color=ACCENT, s=18, zorder=3, edgecolor="none")
    axR.set_yscale("log")
    axR.set_xlabel("horizon, n (steps)")
    axR.set_ylabel("expected total work (step-executions)")
    axR.set_xlim(1, 400)
    axR.legend(loc="best")
    style(axR)

    save(fig, "orchestration-durability.svg")

if __name__ == "__main__":
    means = fig_scheduling()
    fig_durability()
    print("scheduling means (random, spt, hc, c-mu):",
          ", ".join(f"{m:.2f}" for m in means))
    print("done.")
