# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Key-bounded against time-bounded state.

Two operators over the same stream: a per-user counter (state bounded by the
number of keys) and an interval-join buffer that keeps each record for 600
simulated seconds (state bounded by rate × retention). Memory and entries at
two event rates, then the buffer with its eviction removed.
"""
import argparse, json, random, tracemalloc
from collections import deque

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rates", default="100,200", help="Events per simulated second.")
p.add_argument("--seconds", type=int, default=1200, help="Simulated seconds.")
p.add_argument("--users", type=int, default=100_000, help="Distinct users.")
p.add_argument("--interval", type=int, default=600, help="Join retention (s).")
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()

def run(rate, evict=True):
    random.seed(args.seed)
    tracemalloc.start(); counts, buf = {}, deque()
    for t in range(args.seconds):
        for _ in range(rate):
            u = random.randrange(args.users); counts[u] = counts.get(u, 0) + 1; buf.append((t, u))
        if evict:
            while buf and buf[0][0] <= t - args.interval: buf.popleft()         # eviction = the retention rule
    cur, _ = tracemalloc.get_traced_memory(); tracemalloc.stop()
    return {"rate": rate, "eviction": "yes" if evict else "no", "keys": len(counts), "buffer": len(buf), "memory_mb": cur / 1e6}
rates = [int(x) for x in args.rates.split(",")]
rows = [run(r) for r in rates] + [run(rates[-1], evict=False)]
for r in rows: print(f"rate {r['rate']:>4}/s, eviction {r['eviction']:<3}: keys {r['keys']:>7,}  buffered records {r['buffer']:>8,}  memory {r['memory_mb']:6.1f} MB")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"runs": rows}}, f)