# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Sliding windows with panes, and floating-point drift.

(1) A 24-hour sliding count per key refreshed every minute, kept two ways:
each record added to all 1,440 windows it belongs to, or folded once into a
1-minute pane with 1,440 panes summed at emit time. Work and memory for
100,000 records. (2) A running sliding sum of floats kept by add-new,
subtract-old, against the exact sum (math.fsum) after 10 million updates.
"""
import argparse, json, math, random, time, tracemalloc
from collections import deque, defaultdict

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=100_000, help="Records for part 1.")
p.add_argument("--keys", type=int, default=100, help="Keys for part 1.")
p.add_argument("--updates", type=int, default=10_000_000, help="Updates for part 2.")
p.add_argument("--window", type=int, default=1000, help="Part 2: values in the sliding window.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)

W = 1440
recs = [(random.randrange(args.keys), random.randrange(3 * W)) for _ in range(args.records)]     # (key, minute)
def per_window():
    tracemalloc.start(); t = time.perf_counter(); win = defaultdict(int); ops = 0
    for k, m in recs:
        for w in range(m, m + W): win[(k, w)] += 1; ops += 1                  # the record lives in 1,440 windows
    el = time.perf_counter() - t; mem = tracemalloc.get_traced_memory()[0]; tracemalloc.stop()
    return el, mem, ops, win
def panes():
    tracemalloc.start(); t = time.perf_counter(); pane = defaultdict(int); ops = 0
    for k, m in recs: pane[(k, m)] += 1; ops += 1                            # fold once into its 1-minute pane
    out = {}
    for k in range(args.keys):
        for w in range(W - 1, 3 * W, 60):                                     # emit one window per hour, from 1,440 panes
            out[(k, w)] = sum(pane.get((k, m), 0) for m in range(w - W + 1, w + 1)); ops += W
    el = time.perf_counter() - t; mem = tracemalloc.get_traced_memory()[0]; tracemalloc.stop()
    return el, mem, ops, out
a_t, a_m, a_ops, a = per_window(); b_t, b_m, b_ops, b = panes()
same = all(a.get(k, 0) == v for k, v in b.items())
part1 = [{"method": "record in every window", "seconds": a_t, "memory_mb": a_m / 1e6, "per_record_ops": a_ops / args.records},
         {"method": "1-minute panes", "seconds": b_t, "memory_mb": b_m / 1e6, "per_record_ops": 1}]
xs = deque(); running = 0.0; exact = 0.0; checkpoints = []
for i in range(args.updates):
    v = random.uniform(-1, 1) * 10 ** random.randint(-3, 9)
    xs.append(v); running += v
    if len(xs) > args.window: running -= xs.popleft()                         # add-new, subtract-old
    if (i + 1) % (args.updates // 10) == 0:
        exact = math.fsum(xs); checkpoints.append({"updates": i + 1, "running": running, "exact": exact, "drift": running - exact})
summary = {"same_counts": "yes" if same else "no", "final_drift": checkpoints[-1]["drift"], "final_exact": checkpoints[-1]["exact"],
           "max_abs_drift": max(abs(c["drift"]) for c in checkpoints)}
for r in part1: print(f"{r['method']:<24} {r['seconds']:6.2f} s  {r['memory_mb']:7.1f} MB  {r['per_record_ops']:6.0f} updates per record")
print(f"same counts at the emitted windows: {summary['same_counts']}")
for c in checkpoints[-3:]: print(f"after {c['updates']:>10,} updates: running sum {c['running']:.6f}, exact {c['exact']:.6f}, drift {c['drift']:.3g}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"panes": part1, "drift": checkpoints}}, f)