# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "datasketches"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The cardinality wall, and Gorilla encoding.

(1) Series counts for labels region (5), city (50) and courier_tier (3), then
with a user_id label of 100,000 values, at 3 KB per series. (2) A million
timestamps at a 1 s interval with small jitter and a million slowly drifting
values: delta-of-delta zeros, XOR leading/trailing zero bits, the estimated
encoded size against 16 raw bytes per point and zlib; then the same with
noisy values. (3) Minute rollups: the hourly mean from minute (sum, count),
the hourly p99 from minute p99s, and from merged t-digests.
"""
import argparse, json, struct, zlib
import numpy as np
from datasketches import tdigest_double

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--points", type=int, default=1_000_000, help="Points.")
p.add_argument("--users", type=int, default=100_000, help="Values of user_id.")
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()

rng = np.random.default_rng(args.seed)
base_series = 5 * 50 * 3; with_user = base_series * args.users
card = [{"labels": "region, city, courier_tier", "series": base_series, "memory_gb": base_series * 3 / 2**20},
        {"labels": "+ user_id", "series": with_user, "memory_gb": with_user * 3 / 2**20}]
def encoded_bits(ts, vals):
    dod = np.diff(np.diff(ts)); zero = int((dod == 0).sum())
    ts_bits = zero * 1 + (len(dod) - zero) * 14                                 # 1 bit per zero, ~9-14 bits otherwise (upper end)
    b = vals.astype(np.float64).view(np.uint64); x = b[1:] ^ b[:-1]; same = int((x == 0).sum())
    nz = x[x != 0]; lead = np.array([64 - int(v).bit_length() for v in nz.tolist()]); trail = np.array([(int(v) & -int(v)).bit_length() - 1 for v in nz.tolist()])
    meaningful = 64 - lead - trail
    val_bits = same * 1 + int((meaningful + 13).sum())                           # control bits + 6-bit lead + 6-bit length + meaningful bits
    return ts_bits, val_bits, zero / len(dod), same / len(x), float(lead.mean()) if len(lead) else 64, float(trail.mean()) if len(trail) else 64
rows = []
ts = (1_700_000_000_000 + np.arange(args.points) * 1000 + np.where(rng.random(args.points) < 0.05, rng.integers(-3, 4, args.points), 0)).astype(np.int64)
for label, vals in (("slowly drifting values", np.round(20 + np.cumsum(rng.normal(0, 0.01, args.points)), 2)),
                    ("noisy values", 20 + rng.normal(0, 1, args.points))):
    tb, vb, zero, same, lead, trail = encoded_bits(ts, vals)
    raw = args.points * 16; enc = (tb + vb) / 8
    z = len(zlib.compress(ts.tobytes() + vals.astype(np.float64).tobytes(), 6))
    rows.append({"data": label, "dod_zero_pct": zero * 100, "xor_equal_pct": same * 100, "mean_leading_zeros": lead, "mean_trailing_zeros": trail,
                 "raw_mb": raw / 1e6, "gorilla_mb": enc / 1e6, "gorilla_ratio": raw / enc, "zlib_ratio": raw / z})
lat = rng.lognormal(3, 0.8, 3600 * 200)                                           # an hour of latencies, 200 per second
minutes = lat.reshape(60, -1)
true_mean, true_p99 = float(lat.mean()), float(np.percentile(lat, 99))
mean_from_sums = float(minutes.sum(1).sum() / minutes.size); p99_from_p99s = float(np.mean(np.percentile(minutes, 99, axis=1)))
merged = tdigest_double(200)
for m in minutes:
    td = tdigest_double(200)
    for v in m.tolist(): td.update(v)
    merged.merge(td)
rollup = [{"statistic": "hourly mean from minute (sum, count)", "value": mean_from_sums, "true": true_mean},
          {"statistic": "hourly p99 as mean of minute p99s", "value": p99_from_p99s, "true": true_p99},
          {"statistic": "hourly p99 from merged t-digests", "value": float(merged.get_quantile(0.99)), "true": true_p99}]
for r in rollup: r["err_pct"] = (r["value"] / r["true"] - 1) * 100
for r in card: print(f"{r['labels']:<30} {r['series']:>12,} series  ~{r['memory_gb']:8.2f} GB at 3 KB each")
for r in rows: print(f"{r['data']:<24} dod zeros {r['dod_zero_pct']:5.1f}%  XOR equal {r['xor_equal_pct']:5.1f}%  Gorilla ~{r['gorilla_ratio']:4.1f}x  zlib {r['zlib_ratio']:4.1f}x")
for r in rollup: print(f"{r['statistic']:<40} {r['value']:8.2f} (true {r['true']:.2f}, {r['err_pct']:+.2f}%)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"cardinality": card, "encoding": rows, "rollup": rollup}}, f)