usps.py

"""
Download post offices from OpenStreetMap via Overpass, state-by-state
- Resolves states by ISO3166-2 inside Overpass (no Nominatim, avoids Washington/DC mixups).
- Writes one GeoPackage per state (points; EPSG:4326), optional national merge.
- Optional USPS-only filter (operator/brand ~ USPS) or "all post_office" for max speed.

Install:
  pip install requests geopandas shapely pandas tqdm

License & attribution:
  Data © OpenStreetMap contributors, ODbL 1.0.
"""

import os
import json
import time
import argparse
import logging
from typing import Dict, List, Tuple, Optional
import requests
import pandas as pd
import geopandas as gpd
from shapely.geometry import Point
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed

# -------------------- Defaults & Constants --------------------

USER_AGENT = "USPS-OSM-StateDownloader/1.1 (contact: brian.almdale@gmail.com)"
DEFAULT_OVERPASS_ENDPOINTS = [
    "https://overpass-api.de/api/interpreter",
    "https://overpass.kumi.systems/api/interpreter",
    "https://overpass.openstreetmap.fr/api/interpreter",
]

STATES = [
    ("AL","Alabama"), ("AK","Alaska"), ("AZ","Arizona"), ("AR","Arkansas"),
    ("CA","California"), ("CO","Colorado"), ("CT","Connecticut"), ("DE","Delaware"),
    ("DC","District of Columbia"),
    ("FL","Florida"), ("GA","Georgia"), ("HI","Hawaii"), ("ID","Idaho"),
    ("IL","Illinois"), ("IN","Indiana"), ("IA","Iowa"), ("KS","Kansas"),
    ("KY","Kentucky"), ("LA","Louisiana"), ("ME","Maine"), ("MD","Maryland"),
    ("MA","Massachusetts"), ("MI","Michigan"), ("MN","Minnesota"), ("MS","Mississippi"),
    ("MO","Missouri"), ("MT","Montana"), ("NE","Nebraska"), ("NV","Nevada"),
    ("NH","New Hampshire"), ("NJ","New Jersey"), ("NM","New Mexico"), ("NY","New York"),
    ("NC","North Carolina"), ("ND","North Dakota"), ("OH","Ohio"), ("OK","Oklahoma"),
    ("OR","Oregon"), ("PA","Pennsylvania"), ("RI","Rhode Island"), ("SC","South Carolina"),
    ("SD","South Dakota"), ("TN","Tennessee"), ("TX","Texas"), ("UT","Utah"),
    ("VT","Vermont"), ("VA","Virginia"), ("WA","Washington"), ("WV","West Virginia"),
    ("WI","Wisconsin"), ("WY","Wyoming")
]

TERRITORIES = [
    ("PR","Puerto Rico"), ("GU","Guam"), ("VI","United States Virgin Islands"),
    ("AS","American Samoa"), ("MP","Northern Mariana Islands")
]

ISO2_BY_ABBR = {
    "AL":"US-AL","AK":"US-AK","AZ":"US-AZ","AR":"US-AR","CA":"US-CA","CO":"US-CO","CT":"US-CT","DE":"US-DE",
    "DC":"US-DC","FL":"US-FL","GA":"US-GA","HI":"US-HI","ID":"US-ID","IL":"US-IL","IN":"US-IN","IA":"US-IA",
    "KS":"US-KS","KY":"US-KY","LA":"US-LA","ME":"US-ME","MD":"US-MD","MA":"US-MA","MI":"US-MI","MN":"US-MN",
    "MS":"US-MS","MO":"US-MO","MT":"US-MT","NE":"US-NE","NV":"US-NV","NH":"US-NH","NJ":"US-NJ","NM":"US-NM",
    "NY":"US-NY","NC":"US-NC","ND":"US-ND","OH":"US-OH","OK":"US-OK","OR":"US-OR","PA":"US-PA","RI":"US-RI",
    "SC":"US-SC","SD":"US-SD","TN":"US-TN","TX":"US-TX","UT":"US-UT","VT":"US-VT","VA":"US-VA","WA":"US-WA",
    "WV":"US-WV","WI":"US-WI","WY":"US-WY",
    "PR":"US-PR","GU":"US-GU","VI":"US-VI","AS":"US-AS","MP":"US-MP"
}

# -------------------------------------------------------------


def setup_logging(verbose: bool):
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)


def build_overpass_query_iso(iso_code: str, strict_usps: bool, timeout: int) -> str:
    """
    Resolve a state's area by ISO3166-2 inside Overpass, then query post_office features in that area.
    """
    if strict_usps:
        filter_block = '''
  node["amenity"="post_office"]["operator"~"United States Postal Service|USPS"](area.a);
  way["amenity"="post_office"]["operator"~"United States Postal Service|USPS"](area.a);
  relation["amenity"="post_office"]["operator"~"United States Postal Service|USPS"](area.a);

  node["amenity"="post_office"]["brand"~"United States Postal Service|USPS"](area.a);
  way["amenity"="post_office"]["brand"~"United States Postal Service|USPS"](area.a);
  relation["amenity"="post_office"]["brand"~"United States Postal Service|USPS"](area.a);
'''
    else:
        filter_block = '''
  node["amenity"="post_office"](area.a);
  way["amenity"="post_office"](area.a);
  relation["amenity"="post_office"](area.a);
'''

    return f"""
[out:json][timeout:{timeout}];
rel["boundary"="administrative"]["ISO3166-2"="{iso_code}"];
map_to_area->.a;
(
{filter_block}
);
out center tags;
"""


def overpass_request(query: str, endpoints: List[str], max_retries: int = 6) -> dict:
    headers = {"User-Agent": USER_AGENT, "Accept-Encoding": "gzip, deflate"}
    delay = 2.0
    last_err = None

    for attempt in range(max_retries):
        for url in endpoints:
            try:
                r = requests.post(url, data={"data": query}, headers=headers, timeout=120)
                if r.status_code == 200:
                    return r.json()
                if r.status_code in (429, 500, 502, 503, 504):
                    last_err = RuntimeError(f"{url} {r.status_code}: {r.text[:200]}")
                    time.sleep(delay)
                    continue
                r.raise_for_status()
            except requests.RequestException as e:
                last_err = e
                time.sleep(delay)
        delay = min(30.0, delay * 1.7)

    raise RuntimeError(f"Overpass request failed after retries. Last error: {last_err}")


def elements_to_df(elements: List[dict]) -> pd.DataFrame:
    rows = []
    for el in elements:
        et = el.get("type")
        eid = el.get("id")
        tags = el.get("tags", {}) or {}

        if et == "node":
            lat = el.get("lat")
            lon = el.get("lon")
        else:
            c = el.get("center") or {}
            lat = c.get("lat")
            lon = c.get("lon")

        if lat is None or lon is None:
            continue

        rows.append({
            "osm_id": f"{et}/{eid}",
            "osm_type": et,
            "osm_numeric_id": eid,
            "name": tags.get("name"),
            "operator": tags.get("operator"),
            "brand": tags.get("brand"),
            "ref_usps": tags.get("ref:usps") or tags.get("usps:id"),
            "addr_housenumber": tags.get("addr:housenumber"),
            "addr_street": tags.get("addr:street"),
            "addr_unit": tags.get("addr:unit"),
            "addr_city": tags.get("addr:city"),
            "addr_state": tags.get("addr:state"),
            "addr_postcode": tags.get("addr:postcode"),
            "phone": tags.get("phone"),
            "website": tags.get("website"),
            "opening_hours": tags.get("opening_hours"),
            "source_tag": tags.get("source"),
            "lon": float(lon),
            "lat": float(lat),
            "all_tags_json": json.dumps(tags, ensure_ascii=False)
        })
    df = pd.DataFrame(rows)
    if not df.empty:
        df = df.drop_duplicates(subset=["osm_id"]).reset_index(drop=True)
    return df


def write_gpkg_points(df: pd.DataFrame, out_path: str, layer: str = "points") -> Optional[gpd.GeoDataFrame]:
    if df is None or df.empty:
        return None
    gdf = gpd.GeoDataFrame(
        df,
        geometry=[Point(xy) for xy in zip(df["lon"], df["lat"])],
        crs="EPSG:4326"
    )
    # Overwrite file if it exists (single-layer output per state)
    if os.path.exists(out_path):
        os.remove(out_path)
    gdf.to_file(out_path, layer=layer, driver="GPKG")
    return gdf


def run_state(
    abbr: str,
    name: str,
    out_dir: str,
    strict_usps: bool,
    timeout: int,
    endpoints: List[str],
    skip_existing: bool = False,
    quiet: bool = False
) -> Optional[str]:
    gpkg = os.path.join(out_dir, f"usps_{abbr}.gpkg")
    if skip_existing and os.path.exists(gpkg):
        if not quiet:
            logging.info(f"[{abbr}] exists, skipping → {gpkg}")
        return gpkg

    iso = ISO2_BY_ABBR.get(abbr.upper())
    if not iso:
        logging.error(f"[{abbr}] missing ISO3166-2 map entry.")
        return None

    query = build_overpass_query_iso(iso, strict_usps, timeout)

    try:
        data = overpass_request(query, endpoints)
    except Exception as e:
        logging.error(f"[{abbr}] Overpass error: {e}")
        return None

    df = elements_to_df(data.get("elements", []))
    if df.empty:
        logging.warning(f"[{abbr}] No features returned.")
        return None

    gdf = write_gpkg_points(df, gpkg, layer="points")
    if gdf is not None:
        if not quiet:
            logging.info(f"[{abbr}] wrote {len(gdf):,} points → {gpkg}")
        return gpkg
    return None


def merge_all(gpkgs: List[str], out_path: str, layer_name: str = "usps_points_us"):
    frames = []
    for p in gpkgs:
        try:
            g = gpd.read_file(p, layer="points")
            if not g.empty:
                frames.append(g)
        except Exception as e:
            logging.warning(f"Skip merge {p}: {e}")
    if not frames:
        logging.warning("Nothing to merge.")
        return
    merged = pd.concat(frames, ignore_index=True)
    gdf = gpd.GeoDataFrame(merged, geometry="geometry", crs="EPSG:4326")
    if os.path.exists(out_path):
        os.remove(out_path)
    gdf.to_file(out_path, layer=layer_name, driver="GPKG")
    logging.info(f"Merged {len(frames)} layers → {out_path} ({len(gdf):,} points)")


def main():
    ap = argparse.ArgumentParser(description="Download OSM post offices by US state (ISO-based; fast; no Nominatim).")
    ap.add_argument("--out-dir", default="./usps_out", help="Output directory for per-state GPKGs")
    ap.add_argument("--strict-usps", action="store_true",
                    help="Filter to USPS-run facilities (operator/brand ~ USPS). If omitted, grabs ALL post offices.")
    ap.add_argument("--include-territories", action="store_true", help="Add PR, GU, VI, AS, MP")
    ap.add_argument("--merge", action="store_true", help="Merge all state outputs into one GPKG")
    ap.add_argument("--merge-path", default="./usps_points_us.gpkg", help="Path for merged national GPKG")
    ap.add_argument("--timeout", type=int, default=60, help="Overpass timeout per request (seconds)")
    ap.add_argument("--skip-existing", action="store_true", help="Skip states whose GPKG already exists")
    ap.add_argument("--overpass", nargs="*", default=DEFAULT_OVERPASS_ENDPOINTS, help="Overpass endpoints (space-separated)")
    ap.add_argument("--states", default="ALL",
                    help='Comma list like "CA,OR,WA" (default ALL)')
    ap.add_argument("--workers", type=int, default=1,
                    help="Number of parallel states to fetch (1-6 is sensible).")
    ap.add_argument("--verbose", action="store_true", help="Verbose logging")
    args = ap.parse_args()

    setup_logging(args.verbose)
    os.makedirs(args.out_dir, exist_ok=True)

    # Build state list
    target = STATES.copy()
    if args.include_territories:
        target += TERRITORIES

    if args.states != "ALL":
        keep = set([s.strip().upper() for s in args.states.split(",") if s.strip()])
        target = [(a, n) for (a, n) in target if a in keep]
        if not target:
            raise SystemExit("No matching states after --states filtering.")

    done_paths: List[str] = []

    # Process states (optionally in parallel)
    workers = max(1, min(int(args.workers), 12))
    logging.info(f"Fetching {len(target)} states with workers={workers} (strict_usps={args.strict_usps})")

    if workers == 1:
        for abbr, name in tqdm(target, desc="States", unit="state"):
            p = run_state(
                abbr=abbr,
                name=name,
                out_dir=args.out_dir,
                strict_usps=args.strict_usps,
                timeout=args.timeout,
                endpoints=args.overpass,
                skip_existing=args.skip_existing,
                quiet=not args.verbose
            )
            if p:
                done_paths.append(p)
            time.sleep(0.3)  # gentle pacing
    else:
        with ThreadPoolExecutor(max_workers=workers) as ex:
            futs = {
                ex.submit(
                    run_state,
                    abbr, name, args.out_dir, args.strict_usps, args.timeout, args.overpass,
                    args.skip_existing, not args.verbose
                ): (abbr, name)
                for abbr, name in target
            }
            for fut in tqdm(as_completed(futs), total=len(futs), desc="States (parallel)", unit="state"):
                p = fut.result()
                if p:
                    done_paths.append(p)
        # small pause after parallel burst
        time.sleep(0.5)

    if args.merge and done_paths:
        merge_all(done_paths, args.merge_path)

    print("Done. Remember ODbL attribution: © OpenStreetMap contributors.")


if __name__ == "__main__":
    main()