# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""How many keys move when N changes.

100,000 keys on 10 nodes; add an 11th and count the keys that change node
under modulo hashing, a consistent-hash ring with 1 and 128 virtual nodes,
and rendezvous hashing. Then remove a node and see how many successors take
its keys, and time a million lookups.
"""
import argparse, bisect, hashlib, json, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--keys", type=int, default=100_000, help="Keys.")
p.add_argument("--nodes", type=int, default=10, help="Nodes before the change.")
p.add_argument("--lookups", type=int, default=1_000_000, help="Lookups to time.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

h = lambda s: int(hashlib.md5(s.encode()).hexdigest(), 16)
def modulo(keys, nodes):  return {k: nodes[h(k) % len(nodes)] for k in keys}
def rendezvous(keys, nodes): return {k: max(nodes, key=lambda n: h(f"{k}|{n}")) for k in keys}
def ring_fn(nodes, v):
    pts = sorted((h(f"{n}#{i}"), n) for n in nodes for i in range(v)); xs = [x for x, _ in pts]
    return lambda k: pts[bisect.bisect(xs, h(k)) % len(pts)][1]
def ring(keys, nodes, v):
    f = ring_fn(nodes, v); return {k: f(k) for k in keys}

keys = [f"key-{i}" for i in range(args.keys)]
before, after, removed = [f"n{i}" for i in range(args.nodes)], [f"n{i}" for i in range(args.nodes + 1)], [f"n{i}" for i in range(1, args.nodes)]
schemes = {"hash % N": modulo, "ring V=1": lambda k, n: ring(k, n, 1), "ring V=128": lambda k, n: ring(k, n, 128),
           "rendezvous": rendezvous}
rows = []
for name, fn in schemes.items():
    a, b, c = fn(keys, before), fn(keys, after), fn(keys, removed)
    moved = sum(a[k] != b[k] for k in keys) / len(keys)
    load = {}
    for k in keys: load[a[k]] = load.get(a[k], 0) + 1
    heirs = {c[k] for k in keys if a[k] == "n0"}
    rows.append({"scheme": name, "moved_pct": moved * 100, "max_over_min_load": max(load.values()) / min(load.values()),
                 "heirs_of_removed": len(heirs)})
    print(f"{name:<12} add a node: {moved:6.1%} of keys move   load max/min {rows[-1]['max_over_min_load']:5.2f}   "
          f"a removed node's keys go to {len(heirs)} node(s)")
look = [f"q-{i}" for i in range(args.lookups)]
f128 = ring_fn(after, 128)
t = time.perf_counter(); [f128(k) for k in look]; t_ring = time.perf_counter() - t
t = time.perf_counter(); [max(after, key=lambda n: h(f"{k}|{n}")) for k in look[: args.lookups // 10]]; t_rv = (time.perf_counter() - t) * 10
print(f"{args.lookups:,} lookups: ring (binary search) {t_ring:.2f} s   rendezvous (hash every node) {t_rv:.2f} s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"schemes": rows, "lookups": [{"scheme": "ring V=128", "seconds": t_ring}, {"scheme": "rendezvous", "seconds": t_rv}]}}, f)