#!/usr/bin/env nix-shell
#!nix-shell -i "python3 -i" -p python3Packages.shapely python3Packages.fiona python3Packages.requests

print("Starting…")

from collections import defaultdict
import csv
import json
import os
import io

import requests
import fiona
from shapely.geometry import shape, mapping
from shapely.ops import unary_union

print("Completed imports")

# Create lists of zip codes per center and centers per state
print("parsing URC255V.csv…", end="", flush=True)
if not os.path.exists("URC255V.csv"):
    print("Warning: URC255V.csv not found. Please download it from UPS and place it in this directory.")
    exit(1)
# Read URC255V.csv into a dict
reader = csv.DictReader(open("URC255V.csv", encoding="utf-8"))
centers_by_state = defaultdict(set)
zips_by_center = defaultdict(list)
# "CountryCode","PostalLow","PostalHigh","URC25.5V","06/2023"
# "US","55336","55336"," MN 553 0-01"
# "US","55337","55337"," MN 551 9-02"
for row in reader:
    if row["CountryCode"] == "US":
        center = row["URC25.5V"].strip().split('-')[0]
        state = center.split(" ")[0]
        centers_by_state[state].add(center)
        for zip in range(int(row["PostalLow"]), int(row["PostalHigh"]) + 1):
            zips_by_center[center].append(str(zip).zfill(5))
print("complete.")

# Fetch and parse zip code geometries
# TODO: could also get as geopackage or kml file I think? Not sure if either of those is easier to open.
print("fetching zip code data…", end="", flush=True)
res = requests.get("https://www2.census.gov/geo/tiger/GENZ2020/shp/cb_2020_us_zcta520_500k.zip")
res.raise_for_status()
print("complete.")
print("parsing zip code data…", end="", flush=True)
# Yeah, this loads into memory, but the file is only 60ish MiB.
zip_data = list(fiona.io.ZipMemoryFile(io.BytesIO(res.content)).open())
zips = {}
for zip in zip_data: # PIVOT! PIVOT!
    zips[zip.properties["NAME20"]] = zip.geometry
print("complete.")

# Save output
print("writing output files…")
os.makedirs("states", exist_ok=True)
for state, centers in centers_by_state.items():
    print("  "+ state)
    features = []
    for center in centers:
        center_zips = zips_by_center[center]
        features.append({
            "type": "Feature",
            "properties": {
                "center": center,
            },
            "geometry": mapping(unary_union([shape(zips[zip]) for zip in center_zips if zip in zips])),
        })
    with open(f"states/{state}.geojson", "w", encoding="utf-8") as f:
        json.dump(
            {
                "type": "FeatureCollection",
                "features": features,
            },
            f,
        )
print("complete.")