# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "datasketches"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""HyperLogLog: error, merge, and the intersection trap.

An HLL with 2^14 registers over 10 million distinct integers, 20 times; two
sets of 5 million overlapping in 50,000, merged by element-wise max; their
intersection by inclusion-exclusion; then the same with Theta sketches.
"""
import argparse, json
import numpy as np
from datasketches import theta_sketch, theta_intersection, update_theta_sketch

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--p", type=int, default=14, help="log2 of the number of registers.")
p.add_argument("--distinct", type=int, default=10_000_000, help="Distinct items for the error test.")
p.add_argument("--runs", type=int, default=20, help="Runs for the error test.")
p.add_argument("--set-size", type=int, default=5_000_000, help="Items in each of A and B.")
p.add_argument("--overlap", type=int, default=50_000, help="Items in both A and B.")
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()
m = 1 << args.p

def mix(x):
    z = x.astype(np.uint64) + np.uint64(0x9E3779B97F4A7C15)
    z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9)
    z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
    return z ^ (z >> np.uint64(31))

def hll(items):
    h = mix(items); idx = (h >> np.uint64(64 - args.p)).astype(np.int64)
    rest = (h << np.uint64(args.p)) | np.uint64(1 << (args.p - 1))
    rank = (64 - np.floor(np.log2(rest.astype(np.float64))).astype(np.int64))      # leading zeros + 1
    reg = np.zeros(m, dtype=np.int64); np.maximum.at(reg, idx, rank); return reg

def estimate(reg):
    a = 0.7213 / (1 + 1.079 / m); e = a * m * m / np.sum(2.0 ** -reg)
    z = np.sum(reg == 0)
    return m * np.log(m / z) if e <= 2.5 * m and z else e

se = 1.04 / np.sqrt(m)
errs = []
for r in range(args.runs):
    items = np.arange(args.distinct, dtype=np.uint64) + np.uint64(r * 10**9)
    errs.append(estimate(hll(items)) / args.distinct - 1)
within = float(np.mean(np.abs(errs) <= se))
print(f"{args.distinct:,} distinct, {args.runs} runs: mean |error| {np.mean(np.abs(errs)):.3%}, "
      f"{within:.0%} of runs within {se:.2%} (expected about two thirds)")

shared = np.arange(args.overlap, dtype=np.uint64) + np.uint64(5 * 10**10)
A = np.concatenate([np.arange(args.set_size - args.overlap, dtype=np.uint64) + np.uint64(10**10), shared])
B = np.concatenate([np.arange(args.set_size - args.overlap, dtype=np.uint64) + np.uint64(2 * 10**10), shared])
ra, rb = hll(A), hll(B)
merged = np.maximum(ra, rb); union = hll(np.concatenate([A, B]))
same = bool(np.array_equal(merged, union))
inter = estimate(ra) + estimate(rb) - estimate(merged)
print(f"merge by element-wise max equals the HLL of the union: {same}")
print(f"|A ∩ B| by inclusion-exclusion: {inter:,.0f} (true {args.overlap:,}; error {inter / args.overlap - 1:+.0%})")

sa, sb = update_theta_sketch(12), update_theta_sketch(12)
for x in A.tolist(): sa.update(x)
for x in B.tolist(): sb.update(x)
ti = theta_intersection(); ti.update(sa); ti.update(sb)
theta_inter = ti.get_result().get_estimate()
print(f"|A ∩ B| with Theta sketches (k = 4,096): {theta_inter:,.0f} (error {theta_inter / args.overlap - 1:+.0%})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"runs": [{"run": i + 1, "error_pct": e * 100} for i, e in enumerate(errs)],
                   "summary": [{"se_pct": se * 100, "within_pct": within * 100, "mean_abs_err_pct": float(np.mean(np.abs(errs))) * 100,
                                "merge_equal": "yes" if same else "no", "hll_inter": inter, "hll_inter_err_pct": (inter / args.overlap - 1) * 100,
                                "theta_inter": theta_inter, "theta_inter_err_pct": (theta_inter / args.overlap - 1) * 100, "overlap": args.overlap}]}}, f)