# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""FIFO against SRPT on heavy-tailed and steady job sizes.

One server at 90% load, 200,000 jobs, sizes either heavy-tailed (Pareto,
shape 1.8) or almost constant. FIFO against preemptive shortest remaining
processing time (SRPT): mean flow time, p99, and the mean flow time of the
largest 10% of jobs, over several seeds.
"""
import argparse, heapq, json, random, statistics

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--jobs", type=int, default=200_000, help="Jobs per run.")
p.add_argument("--load", type=float, default=0.9, help="Server load.")
p.add_argument("--seeds", type=int, default=5, help="Runs per size distribution.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def gen(n, load, size_fn):
    sizes = [size_fn() for _ in range(n)]
    rate = load / (sum(sizes) / n)                                       # arrival rate that gives this load
    t, jobs = 0.0, []
    for s in sizes:
        t += random.expovariate(rate); jobs.append((t, s))
    return jobs
def fifo(jobs):
    free, out = 0.0, []
    for a, s in jobs:
        start = max(a, free); free = start + s; out.append((s, free - a))  # (size, flow time)
    return out
def srpt(jobs):
    t, i, heap, out, n = 0.0, 0, [], [], len(jobs)
    while i < n or heap:
        if not heap: t = max(t, jobs[i][0])
        while i < n and jobs[i][0] <= t:
            a, s = jobs[i]; heapq.heappush(heap, (s, a, s)); i += 1
        rem, arr, size = heapq.heappop(heap)
        nxt = jobs[i][0] if i < n else float("inf")
        if t + rem <= nxt: t += rem; out.append((size, t - arr))
        else: heapq.heappush(heap, (rem - (nxt - t), arr, size)); t = nxt
    return out
def stats(res):
    res.sort(); flows = sorted(f for _, f in res); big = [f for _, f in res[int(0.9 * len(res)):]]
    return statistics.mean(flows), flows[int(0.99 * len(flows))], statistics.mean(big)
rows = []
for label, fn in (("heavy-tailed", lambda: random.paretovariate(1.8)), ("steady", lambda: random.uniform(0.9, 1.1))):
    for seed in range(args.seeds):
        random.seed(seed); jobs = gen(args.jobs, args.load, fn)
        for pol, f in (("FIFO", fifo), ("SRPT", srpt)):
            m, p99, big = stats(f(jobs)); rows.append({"sizes": label, "seed": seed, "policy": pol, "mean": m, "p99": p99, "big10_mean": big})
summary = []
for label in ("heavy-tailed", "steady"):
    for pol in ("FIFO", "SRPT"):
        rs = [r for r in rows if r["sizes"] == label and r["policy"] == pol]
        summary.append({"sizes": label, "policy": pol, **{k: statistics.median(r[k] for r in rs) for k in ("mean", "p99", "big10_mean")}})
big_better = sum(1 for s in range(args.seeds) if next(r for r in rows if r["sizes"] == "heavy-tailed" and r["seed"] == s and r["policy"] == "SRPT")["big10_mean"]
                 <= next(r for r in rows if r["sizes"] == "heavy-tailed" and r["seed"] == s and r["policy"] == "FIFO")["big10_mean"])
for r in summary: print(f"{r['sizes']:<13} {r['policy']}: mean={r['mean']:8.2f}  p99={r['p99']:9.2f}  largest 10% mean={r['big10_mean']:9.2f}   (median of {args.seeds} seeds)")
print(f"heavy-tailed: SRPT's largest 10% did as well or better than FIFO's in {big_better} of {args.seeds} seeds")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"policies": summary, "seeds": rows, "summary": [{"big_better": big_better, "seeds": args.seeds}]}}, f)