Data pipeline

"""Download + preprocess the NYC TLC data and build the generated instance set.

    python make_data.py            # uses cached raw files if present
    python make_data.py --force    # re-download

Outputs: data/processed/* (gitignored), tests/fixtures/* (small, committed),
instances/*.json (generated orders, committed).
"""
import argparse
import json
import shutil

from delivery import config as C
from delivery import data as D
from delivery.instances import generate_instance


def disk_guard(min_free_mb=150):
    free = shutil.disk_usage(C.ROOT).free / 1e6
    print(f"  free disk: {free:,.0f} MB")
    if free < min_free_mb:
        raise SystemExit(f"stopping: less than {min_free_mb} MB free on the project drive")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    print("[1/5] download"); disk_guard()
    paths = D.download(force=args.force); disk_guard()

    print("[2/5] read + check schema")
    raw = D.load_raw_trips(paths["trips"])
    import pandas as pd
    lookup = pd.read_csv(paths["lookup"])
    cents = D.zone_centroids(paths["shapes"])
    print(f"  {len(raw):,} raw rows, columns OK: {D.NEEDED_COLS}")

    print("[3/5] filter, top zones, sample, split by day")
    trips, zones, stats = D.preprocess(raw, lookup, cents)
    for k, v in stats.items():
        if not isinstance(v, list):
            print(f"  {k:>22}: {v:,}")
    print("  zones:", ", ".join(zones.zone))

    print("[4/5] zone-pair mean distance (train only)")
    dist, cnt, fit = D.pair_distance_matrix(trips, zones)
    stats["pair_distance_fallback"] = fit
    tr = trips[trips.split == "train"]
    g = tr.groupby(["pu", "do", "hour"]).size()
    stats["rare_cells_lt5_train"] = int((g < 5).sum())
    stats["cells_possible"] = len(zones) ** 2 * 24
    stats["cells_seen_train"] = int(len(g))
    night = tr[tr.hour < 6].groupby(["pu", "do"]).size()
    stats["night_pairs_seen_train"] = int(len(night))
    D.save_processed(trips, zones, dist, cnt, stats)
    D.write_fixture(trips, zones, dist)
    print(f"  unseen pairs filled by fit: {fit['unseen_pairs']}; "
          f"(pu,do,hour) cells seen {stats['cells_seen_train']}/{stats['cells_possible']}, "
          f"rare (<5 trips) {stats['rare_cells_lt5_train']}")

    print("[5/5] generated instance set")
    C.INSTANCES.mkdir(exist_ok=True)
    weights = zones.sort_values("idx").n_dropoffs.values
    manifest = []
    for n in (15, 20, 30):
        for seed in range(10):
            inst = generate_instance(n, seed, n_zones=len(zones), zone_weights=weights)
            p = C.INSTANCES / f"{inst.name}.json"
            inst.save(p); manifest.append(p.name)
    (C.INSTANCES / "manifest.json").write_text(json.dumps(manifest, indent=1))
    print(f"  wrote {len(manifest)} instances to instances/")
    disk_guard()


if __name__ == "__main__":
    main()