#!/usr/bin/env python3
"""A hash spreads KEYS evenly but not LOAD when the workload is skewed, DDIA 2e Ch. 7.

16 shards, 10,000 users (user0..user9999). Each user is assigned to a shard by
md5(user_id) % 16 -- a good hash function, so the per-shard KEY COUNTS come out
nearly equal (max/mean ~ 1.0).

But requests are not uniform over users: real workloads follow a Zipfian
popularity law (a celebrity has orders of magnitude more traffic than a typical
user). Weight each user by 1/rank^s with s=1.1 and user0 = the celebrity (rank 1).
Now sum LOAD per shard instead of keys: the single shard holding user0 carries a
hugely disproportionate share of all requests, even though it holds ~1/16 of the
keys. Uniform over keys is not uniform over requests.

The book's application-level fix (§"Skewed Workloads and Relieving Hot Spots"):
split the hot key user0 into 100 sub-keys user0#00..user0#99 and spread its load
evenly across them. Re-hash each sub-key and the WRITE load flattens back toward
~1.x. The catch: a READ of user0 must now gather all 100 sub-keys, hitting every
distinct shard they landed on -- read amplification traded for write spreading.

Pure standard library, deterministic (stable md5 hashing + analytic Zipf weights)
-> the numbers reproduce exactly every run.
"""
import hashlib

SHARDS = 16
USERS = 10_000
ZIPF_S = 1.1          # Zipf exponent; s=1.1 -> a moderately heavy head.
SPLIT = 100           # sub-keys the hot key is split into: user0#00 .. user0#99


def shard_of(key):
    """Assign a key to a shard by a good hash. NOT Python's built-in hash():
    md5 is stable across processes/runs, so these numbers reproduce exactly."""
    return int(hashlib.md5(key.encode()).hexdigest(), 16) % SHARDS


def user_key(i):
    return f"user{i}"


def zipf_weight(i):
    """Popularity weight of user i. user0 is the celebrity (rank 1, heaviest)."""
    rank = i + 1
    return 1.0 / (rank ** ZIPF_S)


def imbalance(loads):
    """max/mean across shards -- 1.0 means perfectly even, higher means hotter."""
    mean = sum(loads) / len(loads)
    return max(loads) / mean, mean


def bar(frac, width=40):
    return "#" * round(frac * width)


def print_shard_table(label, per_shard, as_fraction=False):
    total = sum(per_shard)
    print(label)
    for s in range(SHARDS):
        v = per_shard[s]
        if as_fraction:
            pct = 100.0 * v / total
            print(f"    shard {s:2d} | {pct:5.2f}%  {bar(v / total)}")
        else:
            print(f"    shard {s:2d} | {v:5d}  {bar(v / total)}")


def main():
    # --- Step A: key distribution. Assign every user to a shard, count keys. -----
    key_counts = [0] * SHARDS
    for i in range(USERS):
        key_counts[shard_of(user_key(i))] += 1

    print(f"Step A -- {USERS} keys over {SHARDS} shards, assigned by md5(user_id) % {SHARDS}")
    print_shard_table(f"per-shard KEY COUNT:", key_counts)
    ratio_keys, mean_keys = imbalance(key_counts)
    print(f"    min={min(key_counts)}  max={max(key_counts)}  mean={mean_keys:.1f}")
    print(f"    key-count imbalance  max/mean = {ratio_keys:.2f}   (a good hash spreads keys evenly)\n")

    # --- Step B: request load under a Zipfian popularity law. --------------------
    # Load per shard = sum of the popularity weights of the users it holds.
    load = [0.0] * SHARDS
    for i in range(USERS):
        load[shard_of(user_key(i))] += zipf_weight(i)
    total_load = sum(load)

    hot_user_share = zipf_weight(0) / total_load
    hot_shard = max(range(SHARDS), key=lambda s: load[s])

    print(f"Step B -- SAME shards, but requests follow Zipf (s={ZIPF_S}); user0 = celebrity")
    print_shard_table("per-shard REQUEST LOAD (share of all requests):", load, as_fraction=True)
    ratio_load, mean_load = imbalance(load)
    print(f"    hottest is shard {hot_shard} (holds user0): {100.0 * load[hot_shard] / total_load:.2f}% of ALL requests")
    print(f"    user0 alone is {100.0 * hot_user_share:.2f}% of all requests")
    print(f"    load imbalance  max/mean = {ratio_load:.2f}   (uniform over keys is NOT uniform over load)\n")

    # --- Step C: the fix -- split the hot key across 100 sub-keys. ---------------
    # Remove user0's concentrated load from its shard; redistribute w0/100 to each
    # of user0#00 .. user0#99, re-hashing every sub-key to its own shard.
    user0_shard = shard_of(user_key(0))
    user0_shard_before = load[user0_shard]
    split_load = list(load)
    split_load[user0_shard] -= zipf_weight(0)
    sub_shards = set()
    for d in range(SPLIT):
        sub = f"user0#{d:02d}"
        s = shard_of(sub)
        sub_shards.add(s)
        split_load[s] += zipf_weight(0) / SPLIT

    print(f"Step C -- split the hot key user0 into {SPLIT} sub-keys user0#00..user0#{SPLIT-1:02d}")
    print_shard_table("per-shard WRITE LOAD after splitting the hot key:", split_load, as_fraction=True)
    ratio_split, _ = imbalance(split_load)
    new_hot = max(range(SHARDS), key=lambda s: split_load[s])
    print(f"    user0's shard {user0_shard}: {100.0 * user0_shard_before / total_load:.2f}% -> "
          f"{100.0 * split_load[user0_shard] / total_load:.2f}%  (user0's hot spot is gone)")
    print(f"    but shard {new_hot} -- the NEXT hot key's shard -- now leads at {100.0 * split_load[new_hot] / total_load:.2f}%")
    print(f"    write-load imbalance  max/mean = {ratio_split:.2f}   (was {ratio_load:.2f}: better, not flat --")
    print(f"    Zipf has MANY hot keys; you must split each one, not just the top)\n")

    print(f"    ...and the READ subtlety: a read of user0 must now gather all {SPLIT} sub-keys,")
    print(f"    which landed on {len(sub_shards)} distinct shards -- so a single hot-key read now")
    print(f"    fans out to {len(sub_shards)} of {SHARDS} shards. Write load spread; read load amplified.")


if __name__ == "__main__":
    main()