# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "scikit-learn"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Random split against group split against hash split.

50,000 deliveries for 500 couriers, each with a personal bias and weak
features. A deep decision tree that may memorise courier_id, evaluated on
(a) a random split, (b) GroupKFold by courier and (c) a split where
hash(courier_id) % 10 == 0 goes to test. Then 10% more rows are added and the
test rows that move to training are counted for (a) and (c).
"""
import argparse, hashlib, json
import numpy as np
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import GroupKFold, train_test_split
from sklearn.tree import DecisionTreeRegressor

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=50_000, help="Rows.")
p.add_argument("--couriers", type=int, default=500, help="Couriers.")
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 make(n, start=0):
    cid = rng.integers(0, args.couriers, n); bias = np.random.default_rng(1).normal(0, 5, args.couriers)
    x = rng.normal(0, 1, (n, 3)); y = 30 + bias[cid] + 0.5 * x[:, 0] + rng.normal(0, 3, n)
    return np.column_stack([cid, x]), y, cid, np.arange(start, start + n)
X, y, g, ids = make(args.rows)
model = lambda: DecisionTreeRegressor(min_samples_leaf=5, random_state=0)
def is_test(key): return int(hashlib.md5(str(key).encode()).hexdigest(), 16) % 10 == 0
rows = []
tr, te = train_test_split(np.arange(args.rows), test_size=0.1, random_state=0)
m = model().fit(X[tr], y[tr]); rows.append({"split": "(a) random rows", "test_mae": float(mean_absolute_error(y[te], m.predict(X[te])))})
maes = []
for tri, tei in GroupKFold(n_splits=10).split(X, y, g):
    m = model().fit(X[tri], y[tri]); maes.append(mean_absolute_error(y[tei], m.predict(X[tei])))
rows.append({"split": "(b) GroupKFold by courier", "test_mae": float(np.mean(maes))})
ht = np.array([is_test(c) for c in g]); m = model().fit(X[~ht], y[~ht])
rows.append({"split": "(c) hash(courier_id) % 10 == 0", "test_mae": float(mean_absolute_error(y[ht], m.predict(X[ht])))})
X2, y2, g2, ids2 = make(args.rows // 10, start=args.rows)
all_ids = np.concatenate([ids, ids2]); all_g = np.concatenate([g, g2])
tr2, te2 = train_test_split(np.arange(len(all_ids)), test_size=0.1, random_state=0)
old_test = set(ids[te].tolist()); moved_random = len(old_test - set(all_ids[te2].tolist()))
ht2 = np.array([is_test(c) for c in all_g]); moved_hash = int((ht & ~ht2[:args.rows]).sum())
summary = {"random_moved": moved_random, "random_test_rows": len(old_test), "random_moved_pct": moved_random / len(old_test) * 100,
           "hash_moved": moved_hash, "hash_test_rows": int(ht.sum())}
for r in rows: print(f"{r['split']:<34} test MAE {r['test_mae']:.2f}")
print(f"after adding 10% more rows: random split moved {moved_random:,} of {len(old_test):,} old test rows into training; hash split moved {moved_hash}")

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