#!/usr/bin/env python3
"""Rebuild data/geo/US/*.tsv from the Census Bureau's Gazetteer, population estimates, ZCTA relationships and TIGER/Line files (public domain).

    data-import/geo-us.py [--cache DIR] [--min-streets N] [--place-share F] [--street-share F] [--out DIR]
"""
import collections
import concurrent.futures
import csv
import io
import math
import random
import re
import struct
import sys
import zipfile

import geo
import source
import tsv

GAZETTEER = "https://www2.census.gov/geo/docs/maps-data/data/gazetteer/2026_Gazetteer/2026_Gaz_{}_national.zip"
POPULATION = "https://www2.census.gov/programs-surveys/popest/datasets/2020-2025/{}"
STATES = POPULATION.format("state/totals/NST-EST2025-ALLDATA.csv")
COUNTIES = POPULATION.format("counties/totals/co-est2025-alldata.csv")
PLACES = POPULATION.format("cities/totals/sub-est2025.csv")
ZCTA_PLACE = "https://www2.census.gov/geo/docs/maps-data/data/rel2020/zcta520/tab20_zcta520_place20_natl.txt"
TIGER = "https://www2.census.gov/geo/tiger/TIGER2025/{0}/tl_2025_{1}_{2}.zip"
ESTIMATE = "POPESTIMATE2025"
CDP = "57"
HIGHWAY = re.compile(r"\b(I- |Hwy |Highway |Loop |Rte |Route |Rd )\d")
SUFFIX = re.compile(r" (city and borough|city|town|village|borough|municipality|comunidad|zona urbana|metropolitan government|metro government|consolidated government|unified government|urban county|corporation|plantation)( \(balance\))?$")
# Places whose Census name is a merged government's; the postal city is what an address carries.
NAMES = {"1303440": "Athens", "1304204": "Augusta", "1349008": "Macon", "2148006": "Louisville", "3011397": "Butte", "4732742": "Hartsville", "4752006": "Nashville"}
TIMEZONES = {
    "AK": "America/Anchorage", "AL": "America/Chicago", "AR": "America/Chicago", "AZ": "America/Phoenix",
    "CA": "America/Los_Angeles", "CO": "America/Denver", "CT": "America/New_York", "DC": "America/New_York",
    "DE": "America/New_York", "FL": "America/New_York", "GA": "America/New_York", "HI": "Pacific/Honolulu",
    "IA": "America/Chicago", "ID": "America/Boise", "IL": "America/Chicago", "IN": "America/Indiana/Indianapolis",
    "KS": "America/Chicago", "KY": "America/New_York", "LA": "America/Chicago", "MA": "America/New_York",
    "MD": "America/New_York", "ME": "America/New_York", "MI": "America/Detroit", "MN": "America/Chicago",
    "MO": "America/Chicago", "MS": "America/Chicago", "MT": "America/Denver", "NC": "America/New_York",
    "ND": "America/Chicago", "NE": "America/Chicago", "NH": "America/New_York", "NJ": "America/New_York",
    "NM": "America/Denver", "NV": "America/Los_Angeles", "NY": "America/New_York", "OH": "America/New_York",
    "OK": "America/Chicago", "OR": "America/Los_Angeles", "PA": "America/New_York", "PR": "America/Puerto_Rico",
    "RI": "America/New_York", "SC": "America/New_York", "SD": "America/Chicago", "TN": "America/Chicago",
    "TX": "America/Chicago", "UT": "America/Denver", "VA": "America/New_York", "VT": "America/New_York",
    "WA": "America/Los_Angeles", "WI": "America/Chicago", "WV": "America/New_York", "WY": "America/Denver",
}


def text(data):
    try:
        return data.decode("utf-8-sig")
    except UnicodeDecodeError:
        return data.decode("latin-1")


def gazetteer(cache, kind):
    z = zipfile.ZipFile(io.BytesIO(source.fetch(GAZETTEER.format(kind), cache, f"gaz_{kind}.zip", magic=b"PK")))
    rows = text(z.read(z.namelist()[0])).splitlines()
    header = [h.strip() for h in rows[0].split("|")]
    return [dict(zip(header, (c.strip() for c in row.split("|")))) for row in rows[1:]]


def csv_rows(cache, url, name):
    return list(csv.DictReader(io.StringIO(text(source.fetch(url, cache, name)))))


def dbf_rows(data, wanted):
    """The records of a dBASE file, the wanted fields only."""
    count, header_len, record_len = struct.unpack("<xxxxIHH", data[:12])
    fields, pos = [], 32
    while data[pos] != 0x0D:
        name = data[pos:pos + 11].split(b"\0")[0].decode()
        fields.append((name, data[pos + 16]))
        pos += 32
    pos = header_len
    for _ in range(count):
        record = data[pos:pos + record_len]
        pos += record_len
        if record[:1] == b"*":
            continue
        row, at = {}, 1
        for name, length in fields:
            if name in wanted:
                row[name] = record[at:at + length].decode("utf-8", "replace").strip()
            at += length
        yield row


def tiger_zip(cache, kind, county):
    return source.fetch(TIGER.format(kind.upper(), county, kind), cache, f"tl_{county}_{kind}.zip", magic=b"PK")


def tiger(cache, kind, county, wanted):
    z = zipfile.ZipFile(io.BytesIO(tiger_zip(cache, kind, county)))
    return dbf_rows(z.read(f"tl_2025_{county}_{kind}.dbf"), wanted)


def place_name(geoid, name):
    if geoid in NAMES:
        return NAMES[geoid]
    stripped = SUFFIX.sub("", name)
    if stripped == name:
        print(f"{geoid}: no suffix stripped from {name!r}", file=sys.stderr)
    return stripped


def localities(cache, counties):
    """Each incorporated place with its county, population and centroid."""
    place_population, county_part = {}, collections.defaultdict(list)
    for r in csv_rows(cache, PLACES, "sub-est2025.csv"):
        if r["SUMLEV"] == "162":
            place_population[r["STATE"] + r["PLACE"]] = int(r[ESTIMATE])
        elif r["SUMLEV"] == "157":
            county_part[r["STATE"] + r["PLACE"]].append((int(r[ESTIMATE]), r["STATE"] + r["COUNTY"]))
    out = {}
    for r in gazetteer(cache, "place"):
        geoid, population = r["GEOID"], place_population.get(r["GEOID"], 0)
        if r["FUNCSTAT"] not in ("A", "F", "N") or r["LSAD"] == CDP or not county_part.get(geoid):
            continue
        county = max(county_part[geoid])[1]
        if county not in counties:
            print(f"{geoid} {r['NAME']}: county {county} unknown, dropped", file=sys.stderr)
            continue
        out[geoid] = {"code": geoid, "name": place_name(geoid, r["NAME"]), "municipality": county, "population": population, "lat": r["INTPTLAT"], "lon": r["INTPTLONG"]}
    return out


def drawn(places, share):
    """The share of the places with people, each drawn with odds by its population: Efraimidis and Spirakis's weighted
    sampling without replacement. Each place's random number is seeded by its code, so a place added to or dropped from
    the sources moves only the places near the cut."""
    peopled = [code for code, p in places.items() if p["population"] > 0]
    key = {code: math.log(1 - random.Random(code).random()) / places[code]["population"] for code in peopled}
    return {code: places[code] for code in sorted(peopled, key=lambda c: (key[c], c), reverse=True)[:round(len(peopled) * share)]}


def zcta_localities(cache, localities):
    """Each ZCTA mapped to the incorporated place holding most of its land inside incorporated places, where localities holds that place."""
    parts = {}
    for r in csv.DictReader(io.StringIO(text(source.fetch(ZCTA_PLACE, cache, "zcta-place.txt"))), delimiter="|"):
        if r["GEOID_ZCTA5_20"] and r["GEOID_PLACE_20"] and not r["NAMELSAD_PLACE_20"].endswith(" CDP"):
            parts.setdefault(r["GEOID_ZCTA5_20"], []).append((int(r["AREALAND_PART"]), r["GEOID_PLACE_20"]))
    largest = {zcta: max(p)[1] for zcta, p in parts.items()}
    return {zcta: place for zcta, place in largest.items() if place in localities}


def streets(cache, counties, locality_of_zcta, share, minimum):
    """The names with most TIGER address ranges per place, and the address ranges per ZCTA."""
    with concurrent.futures.ThreadPoolExecutor(3) as pool:
        list(pool.map(lambda c: (tiger_zip(cache, "addr", c), tiger_zip(cache, "featnames", c)), counties))
    addresses, count = collections.Counter(), collections.Counter()
    for county in counties:
        zips = collections.defaultdict(set)
        for r in tiger(cache, "addr", county, {"TLID", "ZIP"}):
            if r["ZIP"] in locality_of_zcta:
                zips[r["TLID"]].add(r["ZIP"])
                addresses[r["ZIP"]] += 1
        for r in tiger(cache, "featnames", county, {"TLID", "FULLNAME", "PAFLAG"}):
            if r["PAFLAG"] == "P" and r["FULLNAME"] and not HIGHWAY.search(r["FULLNAME"]):
                for z in zips.get(r["TLID"], ()):
                    count[(locality_of_zcta[z], r["FULLNAME"])] += 1
    return geo.top_streets(count, share, minimum, "addresses"), addresses


def main():
    a = geo.parser(__doc__, "US", place_share={"type": float, "default": 0.1}).parse_args()
    cache, out = geo.directories(a)

    state_population = {r["STATE"]: r[ESTIMATE] for r in csv_rows(cache, STATES, "nst-est2025.csv") if r["SUMLEV"] == "040"}
    regions = {r["USPS"]: {"abbr": r["USPS"], "code": r["GEOID"], "name": r["NAME"], "population": state_population[r["GEOID"]], "timezone": TIMEZONES[r["USPS"]]} for r in gazetteer(cache, "state")}
    county_population = {r["STATE"] + r["COUNTY"]: r[ESTIMATE] for r in csv_rows(cache, COUNTIES, "co-est2025.csv") if r["SUMLEV"] == "050"}
    counties, unestimated = {}, collections.Counter()
    for r in gazetteer(cache, "counties"):
        if r["GEOID"] not in county_population:
            unestimated[r["USPS"]] += 1
            continue
        counties[r["GEOID"]] = {"code": r["GEOID"], "name": r["NAME"], "region": r["USPS"], "population": county_population[r["GEOID"]]}
    print(f"counties without a population estimate, dropped: {dict(unestimated)}", file=sys.stderr)

    places = localities(cache, counties)
    owner_of_zcta = zcta_localities(cache, places)
    places = drawn({p: places[p] for p in set(owner_of_zcta.values())}, a.place_share)
    locality_of_zcta = {zcta: p for zcta, p in owner_of_zcta.items() if p in places}
    named, addresses = streets(cache, sorted({l["municipality"] for l in places.values()}), locality_of_zcta, a.street_share, a.min_streets)
    zctas = collections.defaultdict(list)
    for zcta, locality in sorted(locality_of_zcta.items()):
        if addresses[zcta]:
            zctas[locality].append(zcta)
    kept, dropped = geo.fence(regions, counties, places, named, zctas)
    geo.log_dropped_localities(dropped["locality"])
    print(f"counties without a place, dropped: {len(dropped['municipality'])}; states without a county, dropped: {sorted(dropped['region'])}", file=sys.stderr)
    places = kept["locality"]

    tsv.write(out / "region.tsv", ["abbr", "code", "name", "population", "timezone"], [r for _, r in sorted(kept["region"].items())])
    tsv.write(out / "municipality.tsv", ["code", "name", "region", "population"], [c for _, c in sorted(kept["municipality"].items())])
    tsv.write(out / "locality.tsv", ["code", "name", "municipality", "population", "lat", "lon"], [l for _, l in sorted(places.items())])
    tsv.write(out / "postal-code.tsv", ["code", "locality", "addresses"], sorted(({"code": z, "locality": l, "addresses": addresses[z]} for l in places for z in zctas[l]), key=lambda r: r["code"]))
    tsv.write(out / "street.tsv", ["name", "locality", "addresses"], [s for locality in sorted(places) for s in named[locality]])


if __name__ == "__main__":
    main()