# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The definition-variant experiment.

"Count active customers in September", answered by three definitions that
each settle the unwritten edge cases differently: refunds, test accounts,
time zone and the period boundary. The chapter's experiment uses three
colleagues; here each colleague is a set of reasonable choices, so the
differences can be counted one by one. Then one shared definition.
"""
import argparse, json
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--customers", type=int, default=50_000, help="Customers.")
p.add_argument("--orders", type=int, default=400_000, help="Orders over three months.")
p.add_argument("--seed", type=float, default=0.42, help="Random seed for DuckDB (between -1 and 1).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

con = duckdb.connect()
con.execute(f"SELECT setseed({args.seed})")
con.execute(f"""CREATE TABLE customers AS SELECT i AS customer_id, random() < 0.02 AS is_test FROM range({args.customers}) t(i)""")
con.execute(f"""CREATE TABLE orders AS SELECT i AS order_id, (random() * {args.customers})::INT AS customer_id,
    TIMESTAMPTZ '2025-08-01 00:00:00+00' + INTERVAL (random() * 92 * 86400) SECOND AS ordered_at_utc,
    random() < 0.04 AS refunded FROM range({args.orders}) t(i)""")
# the three colleagues' reasonable but different choices
variants = {
    "Analyst A: any order, UTC calendar month": """
        SELECT COUNT(DISTINCT customer_id) FROM orders
        WHERE ordered_at_utc >= TIMESTAMPTZ '2025-09-01 00:00:00+00' AND ordered_at_utc < TIMESTAMPTZ '2025-10-01 00:00:00+00'""",
    "Analyst B: exclude refunds and test accounts, UTC": """
        SELECT COUNT(DISTINCT o.customer_id) FROM orders o JOIN customers c USING (customer_id)
        WHERE NOT o.refunded AND NOT c.is_test
          AND ordered_at_utc >= TIMESTAMPTZ '2025-09-01 00:00:00+00' AND ordered_at_utc < TIMESTAMPTZ '2025-10-01 00:00:00+00'""",
    "Analyst C: exclude test accounts, local time (UTC-8), BETWEEN on dates": """
        SELECT COUNT(DISTINCT o.customer_id) FROM orders o JOIN customers c USING (customer_id)
        WHERE NOT c.is_test
          AND CAST(ordered_at_utc - INTERVAL 8 HOUR AS DATE) BETWEEN DATE '2025-09-01' AND DATE '2025-10-01'""",
}
rows = []
for name, sql in variants.items():
    n = con.execute(sql).fetchone()[0]
    rows.append({"definition": name, "active": n}); print(f"{name:<70} {n:>7,}")
diffs = [{"edge_case": "refunded orders", "A": "count", "B": "exclude", "C": "count"},
         {"edge_case": "test accounts", "A": "count", "B": "exclude", "C": "exclude"},
         {"edge_case": "time zone", "A": "UTC", "B": "UTC", "C": "UTC-8"},
         {"edge_case": "end of period", "A": "before 1 Oct", "B": "before 1 Oct", "C": "1 Oct included"}]
spread = max(r["active"] for r in rows) - min(r["active"] for r in rows)
print(f"spread between the three answers: {spread:,} customers ({spread / rows[0]['active']:.1%})")
shared = con.execute(variants["Analyst B: exclude refunds and test accounts, UTC"]).fetchone()[0]
print(f"one written definition, reused by all three: {shared:,}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"definitions": rows, "differences": diffs, "summary": [{"spread": spread,
                   "spread_pct": spread / rows[0]["active"] * 100, "shared": shared}]}}, f)