"""Compute POI proximity counts per postcode from ArcGIS + filtered POIs."""

import argparse
from pathlib import Path

import polars as pl

from pipeline.utils.poi_counts import _count_pois_per_postcode


# POI category groups for proximity counting
POI_GROUPS = {
    "restaurants": ["Restaurant", "Fast Food"],
    "groceries": ["Greengrocer", "Grocery Shop", "Supermarket", "Convenience Store"],
    "parks": ["Park", "Garden", "Nature Reserve"],
    "public_transport": ["Metro or Tram stop", "Rail station", "Bus stop", "Bus station"], # comes from naptan.py
}



def main():
    parser = argparse.ArgumentParser(
        description="Count POIs within radius per postcode"
    )
    parser.add_argument(
        "--arcgis", type=Path, required=True, help="ArcGIS postcode parquet"
    )
    parser.add_argument(
        "--pois", type=Path, required=True, help="Filtered POIs parquet"
    )
    parser.add_argument(
        "--output", type=Path, required=True, help="Output parquet path"
    )
    args = parser.parse_args()

    postcodes = pl.read_parquet(args.arcgis).select(
        pl.col("pcds").alias("postcode"),
        "lat",
        pl.col("long").alias("lon"),
    )

    pois = pl.read_parquet(args.pois)

    result = _count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=2)

    result.write_parquet(args.output)
    size_mb = args.output.stat().st_size / (1024 * 1024)
    print(f"Wrote {args.output} ({size_mb:.1f} MB)")


if __name__ == "__main__":
    main()