# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Recovery-time arithmetic and replay storms.

A queue served at C = 100 jobs/s with arrivals at rho x C; service stops for
60 s, then restarts at full speed: time to drain, against T x rho / (1 - rho).
Then a downstream stage of 120 jobs/s receiving the catch-up plus its own
load; then dropping items older than 30 s.
"""
import argparse, json

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--capacity", type=float, default=100, help="C: jobs per second.")
p.add_argument("--outage", type=float, default=60, help="Seconds without service.")
p.add_argument("--rhos", type=lambda s: [float(v) for v in s.split(",")], default=[0.5, 0.8, 0.9], help="Utilizations.")
p.add_argument("--downstream", type=float, default=120, help="Downstream capacity, jobs per second.")
p.add_argument("--downstream-own", type=float, default=40, help="Downstream's own new load, jobs per second.")
p.add_argument("--max-age", type=float, default=30, help="Drop items older than this, seconds.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
C, T, DT = args.capacity, args.outage, 0.1

def drain(rho, max_age=None):
    from collections import deque
    q, t, lam, drained, down_peak = deque(), 0.0, rho * C, None, 0.0
    carry = 0.0
    while t < 5000:
        carry += lam * DT
        while carry >= 1: q.append(t); carry -= 1
        if max_age is not None:
            while q and t - q[0] > max_age: q.popleft()
        served = 0
        if t >= T:
            for _ in range(int(C * DT)):
                if q: q.popleft(); served += 1
        down_load = served / DT + args.downstream_own
        down_peak = max(down_peak, down_load)
        if t > T and not q and drained is None: drained = t - T
        if drained is not None: break
        t += DT
    return drained, down_peak

rows = []
for rho in args.rhos:
    d, peak = drain(rho)
    rows.append({"rho": rho, "drain_s": d, "formula_s": T * rho / (1 - rho), "downstream_peak": peak,
                 "downstream_overloaded": "yes" if peak > args.downstream else "no"})
    print(f"rho = {rho}: drained after {d:6.1f} s (formula {T * rho / (1 - rho):6.1f} s); downstream peak {peak:.0f}/s vs {args.downstream:.0f}/s capacity")
d, _ = drain(0.8, args.max_age)
print(f"rho = 0.8, dropping items older than {args.max_age:g} s: drained after {d:.1f} s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"rhos": rows, "summary": [{"drop_drain_s": d, "max_age": args.max_age}]}}, f)