# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "scipy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""p-values measure your row count.

A sample with mean 100.1 against a baseline with mean 100 (sd 15), a 0.07%
shift, compared with the two-sample Kolmogorov-Smirnov test and with the
Population Stability Index over 10 quantile buckets, for n from 10^3 to 10^7.
"""
import argparse, json
import numpy as np
from scipy.stats import ks_2samp

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--shift", type=float, default=0.1, help="Shift of the sample's mean.")
p.add_argument("--max-exp", type=int, default=7, help="Largest n is 10^max-exp.")
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)

def psi(a, b, buckets=10):
    edges = np.quantile(a, np.linspace(0, 1, buckets + 1)); edges[0], edges[-1] = -np.inf, np.inf
    pa = np.histogram(a, edges)[0] / len(a); pb = np.histogram(b, edges)[0] / len(b)
    pa, pb = np.clip(pa, 1e-6, None), np.clip(pb, 1e-6, None)
    return float(np.sum((pb - pa) * np.log(pb / pa)))

rows = []
for e in range(3, args.max_exp + 1):
    n = 10 ** e
    a = rng.normal(100, 15, n); b = rng.normal(100 + args.shift, 15, n)
    pv = ks_2samp(a, b).pvalue
    rows.append({"n": n, "p_value": float(max(pv, 1e-300)), "psi": psi(a, b)})
    print(f"n = {n:>10,}   KS p-value {pv:10.3g}   PSI {rows[-1]['psi']:.6f}")

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