#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Geo features cache (flood zone, transit, dissemination-area statistics, CBD distance) for every analytic row.

    python3 scripts/04_geo.py [--synthetic] [--limit N]
→ data/{processed|synthetic}/geo_features.parquet, results/data/geo_summary*.json
Requires data/raw/external/{inondation,transit,quartier}.db (copied from M4M64b:~/apps/immo-ka/data/).
"""
from __future__ import annotations

import argparse
import json
import time

import polars as pl

from agentic_valuation.config import PROCESSED, RESULTS, SYNTHETIC
from agentic_valuation.features.geo import GeoLayers, geo_features_batch


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--synthetic", action="store_true")
    ap.add_argument("--limit", type=int, default=None, help="only the first N rows (smoke test)")
    ap.add_argument("--chunk", type=int, default=100_000)
    args = ap.parse_args()
    base = SYNTHETIC if args.synthetic else PROCESSED
    if not GeoLayers.available():
        raise SystemExit("external layers missing: scp M4M64b:~/apps/immo-ka/data/{inondation,transit,quartier}.db data/raw/external/")
    t0 = time.perf_counter()
    layers = GeoLayers.load()
    t_load = time.perf_counter() - t0
    analytic = pl.read_parquet(base / "analytic.parquet", columns=["tx_id", "lat", "lng", "x_m", "y_m"])
    if args.limit:
        analytic = analytic.head(args.limit)
    parts = []
    for s0 in range(0, analytic.height, args.chunk):
        parts.append(geo_features_batch(analytic.slice(s0, args.chunk), layers))
        print(f"  {min(s0 + args.chunk, analytic.height)}/{analytic.height} rows, {time.perf_counter() - t0:.0f} s", flush=True)
    geo = pl.concat(parts)
    out = base / "geo_features.parquet"
    geo.write_parquet(out, compression="zstd")
    summary = {
        "synthetic": args.synthetic,
        "rows": geo.height,
        "layers_load_seconds": round(t_load, 1),
        "seconds": round(time.perf_counter() - t0, 1),
        "n_flood_polygons": len(layers.flood_geoms),
        "n_coverage_polygons": len(layers.coverage_geoms),
        "n_dissemination_areas": len(layers.da_geoms),
        "n_transit_stops": layers.n_stops,
        "share_flood_mapped": float(geo["flood_mapped"].mean()),
        "share_in_flood_zone_among_mapped": float(geo.filter(pl.col("flood_mapped"))["in_flood_zone"].mean()) if geo["flood_mapped"].any() else None,
        "flood_zone_kinds": {k: int(v) for k, v in geo.group_by("flood_zone_kind").len().iter_rows() if k is not None},
        "share_with_da": float(geo["dauid"].is_not_null().mean()),
        "transit_stops_quantiles": {q: float(geo["transit_stops_within_1km"].quantile(q)) for q in (0.5, 0.75, 0.9)},
        "cbd_counts": {k: int(v) for k, v in geo.group_by("cbd_name").len().iter_rows()},
        "missingness_pct": {c: round(100.0 * geo[c].null_count() / geo.height, 3) for c in geo.columns},
        "output": str(out),
    }
    sp = RESULTS / "data" / ("geo_summary_synthetic.json" if args.synthetic else "geo_summary.json")
    sp.parent.mkdir(parents=True, exist_ok=True)
    sp.write_text(json.dumps(summary, indent=1, ensure_ascii=False))
    print(json.dumps({k: summary[k] for k in ("rows", "layers_load_seconds", "seconds", "share_flood_mapped", "share_in_flood_zone_among_mapped", "share_with_da", "cbd_counts")}, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()