#!/usr/bin/env python3
"""DDIA Ch. 7 (Sharding) -- a monotonic partition key makes a key-range hot spot.

The book's sensor-network example. With KEY-RANGE sharding, each shard owns a
CONTIGUOUS range of keys (that is what makes a range scan cheap: adjacent keys
live together). But if the partition key is a timestamp -- monotonically
increasing -- then the maximum key always lives in the LAST shard, so a
write-as-you-go workload funnels 100% of writes into that one shard while every
other shard sits idle. A self-inflicted write hot spot.

The fix the book gives: prefix the key with something high-cardinality (the
sensor id) so the write point scatters across the sensor-id range. Writes spread
evenly -- at a documented cost: a time-range query across ALL sensors, which was
one contiguous scan under the timestamp key, now needs a separate scan per
sensor and touches every shard.

Two partition-key CHOICES over the identical stream of 100,000 readings; the
same ordering that makes range scans cheap is the ordering that makes
sequential-key writes hot.

Deterministic: the only randomness is sensor assignment, seeded with
random.Random(42), so every number below reproduces byte-for-byte. Stdlib only.
"""

import bisect
import random

N_SHARDS = 12
N_READINGS = 100_000
N_SENSORS = 1_000

# --- A simplified calendar: a 360-day year of seconds, one shard per month ---
DAY = 24 * 3600
MONTH = 30 * DAY          # 2,592,000 s
YEAR = N_SHARDS * MONTH   # 31,104,000 s

# The workload arrives "now": the current (last) month. Readings stream in with
# monotonically increasing timestamps, one per second, starting at the first
# second of month 11 (the newest range).
BASE_TS = (N_SHARDS - 1) * MONTH          # start of month 11


def build_workload():
    """The identical stream both schemes shard: (sensor_id, timestamp).

    Timestamps are monotonically increasing (append-only, one per second),
    all inside the current month. Sensors are drawn from N_SENSORS, seeded.
    """
    rng = random.Random(42)
    readings = []
    for i in range(N_READINGS):
        sensor = rng.randrange(N_SENSORS)   # 0..999, seeded
        ts = BASE_TS + i                    # strictly increasing
        readings.append((sensor, ts))
    return readings


def sid(sensor):
    """Zero-padded, fixed-width sensor id -- so lexicographic == numeric order."""
    return f"s{sensor:04d}"


# --- Scheme A: range-shard by TIMESTAMP -------------------------------------
# Shard m owns the contiguous time range [m*MONTH, (m+1)*MONTH). The partition
# key is the timestamp itself.
def shard_by_timestamp(ts):
    return ts // MONTH


# --- Scheme B: range-shard by "sensor#timestamp" ----------------------------
# The partition key is the string f"{sid}#{ts}", so keys sort by SENSOR first.
# Each shard owns a contiguous slice of the sensor-id space. Boundaries cut the
# 1,000 sensors into 12 contiguous ranges; we route by the sensor-id string.
SENSOR_BOUNDARIES = [sid(round(k * N_SENSORS / N_SHARDS)) for k in range(1, N_SHARDS)]


def shard_by_sensor_key(sensor):
    # number of boundaries the key sorts at-or-after == shard index
    return bisect.bisect_right(SENSOR_BOUNDARIES, sid(sensor))


def bar(count, total, width=40):
    filled = round(width * count / total) if total else 0
    return "#" * filled + "." * (width - filled)


def main():
    readings = build_workload()

    print("=" * 68)
    print("KEY-RANGE SHARDING: a monotonic partition key vs. a scattered one")
    print("=" * 68)
    print(f"{N_READINGS:,} sensor readings from {N_SENSORS:,} sensors, streaming in")
    print(f"with monotonically increasing timestamps (one per second), all inside")
    print(f"the current month. {N_SHARDS} key-range shards.")
    print(f"    timestamp window written: [{BASE_TS:,} .. {BASE_TS + N_READINGS - 1:,}]  "
          f"(all within month {N_SHARDS - 1})")
    print()

    # --- Scenario A ---------------------------------------------------------
    print("-" * 68)
    print("SCENARIO A -- partition key = timestamp (range-shard by month)")
    print("-" * 68)
    print("Shard m owns the contiguous time range [m*MONTH, (m+1)*MONTH):")
    for m in range(N_SHARDS):
        print(f"    shard {m:>2} : [{m * MONTH:>10,} .. {(m + 1) * MONTH:>10,})")
    print()

    counts_a = [0] * N_SHARDS
    for sensor, ts in readings:
        counts_a[shard_by_timestamp(ts)] += 1

    print("per-shard write counts:")
    for m in range(N_SHARDS):
        print(f"    shard {m:>2} : {counts_a[m]:>7,}  [{bar(counts_a[m], N_READINGS)}]")
    busiest_a = max(counts_a)
    idle_a = sum(1 for c in counts_a if c == 0)
    print(f"    busiest shard holds {busiest_a:,} of {N_READINGS:,} writes "
          f"= {100 * busiest_a / N_READINGS:.1f}%")
    print(f"    {idle_a} of {N_SHARDS} shards are completely idle "
          f"-- a self-inflicted write HOT SPOT.")
    print()

    # --- Scenario B ---------------------------------------------------------
    print("-" * 68)
    print("SCENARIO B -- partition key = f\"{sensor}#{timestamp}\" (range-shard by sensor)")
    print("-" * 68)
    print("Same readings, but the key now sorts by SENSOR first. Each shard owns")
    print("a contiguous slice of the 1,000-sensor id space:")
    lo = "s0000"
    for j in range(N_SHARDS):
        hi = SENSOR_BOUNDARIES[j] if j < N_SHARDS - 1 else "s1000"
        print(f"    shard {j:>2} : [{lo} .. {hi})")
        lo = hi
    print()

    counts_b = [0] * N_SHARDS
    for sensor, ts in readings:
        counts_b[shard_by_sensor_key(sensor)] += 1

    print("per-shard write counts:")
    for j in range(N_SHARDS):
        print(f"    shard {j:>2} : {counts_b[j]:>7,}  [{bar(counts_b[j], N_READINGS)}]")
    mx, mn = max(counts_b), min(counts_b)
    mean = N_READINGS / N_SHARDS
    print(f"    max {mx:,}   min {mn:,}   mean {mean:,.0f}")
    print(f"    imbalance: busiest shard = {mx / mean:.3f}x the mean "
          f"(vs {busiest_a / mean:.1f}x under scenario A).")
    print(f"    writes now SPREAD across all {N_SHARDS} shards.")
    print()

    # --- The cost: a cross-sensor time-range query --------------------------
    print("-" * 68)
    print("THE COST -- query: all readings in a time window, ACROSS ALL SENSORS")
    print("-" * 68)
    t0 = BASE_TS + 40_000
    t1 = BASE_TS + 60_000
    matching = [(sensor, ts) for sensor, ts in readings if t0 <= ts <= t1]
    print(f"window [t0,t1] = [{t0:,} .. {t1:,}]  ->  {len(matching):,} readings match")
    print()

    shards_a = {shard_by_timestamp(ts) for _, ts in matching}
    shards_b = {shard_by_sensor_key(sensor) for sensor, _ in matching}
    print(f"scheme A (key = timestamp)      : matching keys are CONTIGUOUS in time")
    print(f"    -> shards touched = {len(shards_a):>2}  {sorted(shards_a)}")
    print(f"scheme B (key = sensor#timestamp): matching keys are SCATTERED by sensor")
    print(f"    -> shards touched = {len(shards_b):>2}  {sorted(shards_b)}")
    print()
    print(f"Spreading the writes cost the query {len(shards_b)}x the shards to scan:")
    print(f"one contiguous range scan became {len(shards_b)} scans (one per shard).")


if __name__ == "__main__":
    main()