# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "pandas", "scikit-learn"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Leakage inflates the score: naive join against as-of join.

20,000 deliveries over 200 days for 100 couriers whose speed drifts. A
feature table holds each courier's 30-day average speed per day, valid from
the end of that day; the label depends slightly on the speed of the day after
the delivery (a deliberate leak). A model trained on a random split of a
naive join (each courier's latest feature value) against one trained on a
temporal split of an as-of join (the value known at delivery time); both are
scored offline and on a later, unseen period.
"""
import argparse, json
import numpy as np, pandas as pd
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.metrics import mean_absolute_error

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
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)
N_COURIERS, N_DAYS, N_EVENTS = 100, 200, 20_000
days = pd.date_range("2025-01-01", periods=N_DAYS, freq="D")
speed = 30 + np.cumsum(rng.normal(0, 0.8, (N_COURIERS, N_DAYS)), axis=1)          # each courier's true daily speed drifts
feats = pd.DataFrame({"courier_id": np.repeat(np.arange(N_COURIERS), N_DAYS), "valid_time": np.tile(days + pd.Timedelta(days=1), N_COURIERS),
                      "avg_speed_30d": pd.DataFrame(speed.T).rolling(30, min_periods=1).mean().T.to_numpy().ravel()})
events = pd.DataFrame({"courier_id": rng.integers(0, N_COURIERS, N_EVENTS),
                       "ts": days[0] + pd.to_timedelta(rng.uniform(0, (N_DAYS - 1) * 24, N_EVENTS), unit="h")})
events["ts"] = events.ts.astype("datetime64[ns]"); feats["valid_time"] = feats.valid_time.astype("datetime64[ns]")
day = (events.ts - days[0]).dt.days.to_numpy()
events["delivery_time"] = 600 / speed[events.courier_id, np.minimum(day + 1, N_DAYS - 1)] + rng.normal(0, 1.0, N_EVENTS)   # the deliberate leak
events = events.sort_values("ts"); feats = feats.sort_values("valid_time")
naive = events.merge(feats.groupby("courier_id").last().reset_index(), on="courier_id")      # today's value
asof = pd.merge_asof(events, feats, left_on="ts", right_on="valid_time", by="courier_id", direction="backward")   # value known at ts
X, y = ["avg_speed_30d"], "delivery_time"
cut, later_cut = days[0] + pd.Timedelta(days=120), days[0] + pd.Timedelta(days=160)
fit = lambda d: GradientBoostingRegressor(random_state=0).fit(d[X], d[y])
mae = lambda m, d: float(mean_absolute_error(d[y], m.predict(d[X])))
later = asof[asof.ts >= later_cut].dropna()                                        # truly unseen, features as served
past = naive[naive.ts < later_cut]; train = past.sample(frac=0.8, random_state=0); m_naive = fit(train)
hist = asof[asof.ts < later_cut].dropna(); m_asof = fit(hist[hist.ts < cut])
rows = [{"model": "naive join, random split", "offline_mae": mae(m_naive, past.drop(train.index)), "later_mae": mae(m_naive, later)},
        {"model": "as-of join, temporal split", "offline_mae": mae(m_asof, hist[hist.ts >= cut]), "later_mae": mae(m_asof, later)}]
for r in rows: r["gap"] = r["later_mae"] - r["offline_mae"]
for r in rows: print(f"{r['model']:<28} offline MAE {r['offline_mae']:.2f}   later period MAE {r['later_mae']:.2f}")

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