# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""How many rows must you sample?

(a) 1,000,000 rows with a defect rate p; for sample sizes from 300 to
300,000, the share of 1,000 trials that find at least one defect, against
1 - (1 - p)^n and the rule n = 3 / p for 95%. (b) 100,000 customers with 3%
postcode typos, blocked by postcode, then by postcode or phone prefix.
"""
import argparse, json, random
from collections import defaultdict
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=1_000_000, help="Rows in the table.")
p.add_argument("--rates", type=lambda s: [float(v) for v in s.split(",")], default=[1e-3, 1e-4, 1e-5], help="Defect rates p.")
p.add_argument("--samples", type=lambda s: [int(v) for v in s.split(",")], default=[300, 3000, 30000, 300000], help="Sample sizes n.")
p.add_argument("--trials", type=int, default=1000, help="Trials per (p, n).")
p.add_argument("--customers", type=int, default=100_000, help="Part (b): customers.")
p.add_argument("--typo", type=float, default=0.03, help="Part (b): share of duplicates with a postcode typo.")
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); random.seed(args.seed)

rows = []
for rate in args.rates:
    defects = rng.random(args.rows) < rate
    for n in args.samples:
        found = np.mean([defects[rng.integers(0, args.rows, n)].any() for _ in range(args.trials)])
        rows.append({"p": rate, "n": n, "found_pct": float(found * 100), "formula_pct": (1 - (1 - rate) ** n) * 100})
    print(f"p = {rate:g}: " + "   ".join(f"n={r['n']:,}: {r['found_pct']:5.1f}% (formula {r['formula_pct']:5.1f}%)"
                                         for r in rows if r["p"] == rate) + f"   rule 3/p = {3 / rate:,.0f}")

# (b) blocking caps recall
people = [{"id": i, "pc": f"{random.randrange(5000):04d}", "phone": f"{random.randrange(1000):03d}"} for i in range(args.customers)]
dups = []
for r in random.sample(people, args.customers // 10):
    pc = r["pc"] if random.random() > args.typo else f"{random.randrange(5000):04d}"
    dups.append({"id": r["id"], "pc": pc, "phone": r["phone"]})
recs = people + dups
plans = {"postcode": [("pc",)], "postcode, then phone prefix": [("pc",), ("phone",)]}
blocking = []
for name, passes in plans.items():
    cand = set()
    for keys in passes:
        groups = defaultdict(list)
        for i, r in enumerate(recs): groups[tuple(r[k] for k in keys)].append(i)
        for g in groups.values():
            for a in range(len(g)):
                for b in range(a + 1, len(g)): cand.add((g[a], g[b]))
    true = sum(1 for a, b in cand if recs[a]["id"] == recs[b]["id"])
    blocking.append({"plan": name, "recall_pct": true / len(dups) * 100, "candidate_pairs": len(cand)})
    print(f"blocking by {name:<28} recall {true / len(dups):6.1%}   candidate pairs {len(cand):,}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"sampling": rows, "blocking": blocking}}, f)