# /// 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
"""Event time against availability time.

The leakage lab's features, each now also carrying available_time =
valid_time + 9 hours (a nightly batch that lands in the morning). An as-of
join on valid_time against one on available_time: how many events receive a
different feature value, against 9/24; then a model trained on each, scored on
a later period with features joined as they would really have been served.
"""
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("--lag-hours", type=float, default=9, help="Hours between valid_time and available_time.")
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)
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()})
feats["valid_time"] = feats.valid_time.astype("datetime64[ns]"); feats["available_time"] = feats.valid_time + pd.Timedelta(hours=args.lag_hours)
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]"); day = (events.ts - days[0]).dt.days.to_numpy()
events["delivery_time"] = 600 / speed[events.courier_id, np.minimum(day, N_DAYS - 1)] + rng.normal(0, 1.0, N_EVENTS)
events = events.sort_values("ts")
by_valid = pd.merge_asof(events, feats.sort_values("valid_time"), left_on="ts", right_on="valid_time", by="courier_id", direction="backward")
by_avail = pd.merge_asof(events, feats.sort_values("available_time").drop(columns="valid_time"), left_on="ts", right_on="available_time", by="courier_id", direction="backward")
both = by_valid[["avg_speed_30d"]].reset_index(drop=True).join(by_avail[["avg_speed_30d"]].reset_index(drop=True), rsuffix="_avail").dropna()
differ = float((both.avg_speed_30d != both.avg_speed_30d_avail).mean())
cut = days[0] + pd.Timedelta(days=150); X, y = ["avg_speed_30d"], "delivery_time"
later = by_avail[by_avail.ts >= cut].dropna()                                      # served features are the available ones
rows = []
for label, d in (("joined on valid_time", by_valid), ("joined on available_time", by_avail)):
    tr = d[d.ts < cut].dropna(); te = d[d.ts >= cut].dropna()
    m = GradientBoostingRegressor(random_state=0).fit(tr[X], tr[y])
    rows.append({"training": label, "offline_mae": float(mean_absolute_error(te[y], m.predict(te[X]))), "served_mae": float(mean_absolute_error(later[y], m.predict(later[X])))})
summary = {"differ_pct": differ * 100, "expected_pct": args.lag_hours / 24 * 100}
print(f"events whose feature value differs between the two joins: {differ:.1%} (expected {args.lag_hours / 24:.1%})")
for r in rows: print(f"{r['training']:<26} offline MAE {r['offline_mae']:.3f}   MAE with features as served {r['served_mae']:.3f}")

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