# /// script
# requires-python = ">=3.10"
# dependencies = ["sortedcontainers"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""SUM versus MAX under deletes.

An incremental SUM and an incremental MAX per group, fed the same stream of
(group, value, +1/-1) changes. SUM needs nothing but its running total; MAX
has no inverse, so deleting the current maximum forces a rescan of the group,
or a stored multiset of every value.
"""
import argparse, json, random, tracemalloc
from collections import defaultdict
from sortedcontainers import SortedList

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--changes", type=int, default=200_000, help="Changes in the stream.")
p.add_argument("--groups", type=int, default=100, help="Groups (for example cities).")
p.add_argument("--delete-share", type=float, default=0.3, help="Share of changes that delete an existing value.")
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)

# The stream: inserts of new values, and deletes of values inserted earlier
live, stream = defaultdict(list), []
for _ in range(args.changes):
    g = random.randrange(args.groups)
    if live[g] and random.random() < args.delete_share:
        v = live[g].pop(random.randrange(len(live[g]))); stream.append((g, v, -1))
    else:
        v = round(random.uniform(1, 1000), 2); live[g].append(v); stream.append((g, v, +1))


def run_sum():
    sums = defaultdict(float)
    for g, v, sign in stream:
        sums[g] += sign * v                            # O(1) for insert and delete
    return sums

def run_max_rescan():
    """MAX with the group's values kept as a plain list; a delete of the max rescans the group."""
    maxes, values, rescans, scanned = {}, defaultdict(list), 0, 0
    for g, v, sign in stream:
        if sign > 0:
            values[g].append(v); maxes[g] = v if maxes.get(g) is None else max(maxes[g], v)
        else:
            values[g].remove(v)
            if v == maxes[g]:                          # the max was deleted: nothing in it says what is next
                rescans += 1; scanned += len(values[g])
                maxes[g] = max(values[g]) if values[g] else None
    return maxes, rescans, scanned

def run_max_sorted():
    """MAX with a sorted multiset per group: no rescans, but every live value is stored."""
    values = defaultdict(SortedList)
    for g, v, sign in stream:
        if sign > 0: values[g].add(v)
        else: values[g].remove(v)
    return {g: (s[-1] if s else None) for g, s in values.items()}


def measured(fn):
    tracemalloc.start(); out = fn(); _, peak = tracemalloc.get_traced_memory(); tracemalloc.stop()
    return out, peak

sums, sum_mem = measured(run_sum)
(maxes, rescans, scanned), rescan_mem = measured(run_max_rescan)
sorted_max, sorted_mem = measured(run_max_sorted)
deletes = sum(1 for _, _, s in stream if s < 0)
assert all(maxes[g] == sorted_max[g] for g in maxes)

print(f"{len(stream):,} changes over {args.groups} groups, {deletes:,} of them deletes")
print(f"SUM:                  0 rescans, state = {args.groups} running totals, peak memory {sum_mem / 1e3:,.0f} KB")
print(f"MAX, rescan on delete: {rescans:,} rescans reading {scanned:,} values, peak memory {rescan_mem / 1e6:,.1f} MB")
print(f"MAX, sorted multiset:  0 rescans, every live value stored, peak memory {sorted_mem / 1e6:,.1f} MB")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "approaches": [{"approach": "SUM", "rescans": 0, "values_read": 0, "peak_mb": sum_mem / 1e6},
                           {"approach": "MAX, rescan on delete", "rescans": rescans, "values_read": scanned, "peak_mb": rescan_mem / 1e6},
                           {"approach": "MAX, sorted multiset", "rescans": 0, "values_read": 0, "peak_mb": sorted_mem / 1e6}],
            "summary": [{"deletes": deletes, "rescans": rescans, "scanned": scanned, "sum_kb": sum_mem / 1e3,
                         "rescan_mb": rescan_mem / 1e6, "sorted_mb": sorted_mem / 1e6}]}}, f)