"""Journey-to-work microdata for New York City, 1990 to 2014, cut to one small table per year.

    py src/extract_commute_history.py                 # every year whose raw file is on disk
    py src/extract_commute_history.py --years 2005 2010

The 2015-2024 one-year files are already extracted for the study area by
extract_metro_pums.py --acs1; this script covers the years before that, which arrive in
three formats:

* ACS 2005-2014 one-year PUMS (csv_pny.zip): CSV, the same variables as the modern files
  under their pre-2019 names (JWTR for means of transportation, ST for state). 2005-2011
  carry 2000-vintage PUMAs, 2012-2014 the 2010 vintage. Person replicate weights are
  lower-case pwgtp1..80.
* Census 2000 5% PUMS (REVISEDPUMS5_36.TXT): fixed-width housing (H) and person (P)
  records, layout from the Bureau's record-layout workbook. Travel time is TRVTIME,
  means of transportation TRVMNS, place of work POWSTATE/POWPUMA1, the person weight
  PWEIGHT. No replicate weights; the article treats the 2000 point as a large sample
  (5% of the population) and reports it without a margin.
* Census 1990 PUMS A (PUMSAXNY.TXT): fixed-width, layout from PUMSUSDD, variables
  TRAVTIME, MEANS, POWSTATE, POWPUMA, PWGT1.

Every year is written to data/interim/acs1/commute_<year>.csv.gz with one canonical set
of columns:

    year, puma, borough, weight, age, sex, esr (1 employed at work, 2 employed not at
    work, 3 unemployed, 6 not in labour force; 1990/2000 coded to the same scheme),
    mode (JWTRNS 2019 coding: 1 car, 2 bus, 3 subway, 4 commuter rail, 5 light rail,
    6 ferry, 7 taxi, 8 motorcycle, 9 bicycle, 10 walked, 11 worked from home, 12 other),
    minutes (one-way, NaN if not a commuter), pow_state, pow_puma, pow_manhattan (place
    of work is Manhattan, by that year's place-of-work PUMA code), earnings (own
    earnings in that year's dollars), hours, race (White/Black/Asian/Hispanic/Other),
    foreign_born, ba_plus, occ (census occupation code, string), depart (departure time
    code as published, where present), ind (census industry code of that year's
    scheme, string), naics (the NAICS-based code where the file carries one, 2000 and
    2005-2014), cow (class of worker as published: 1990 CLASS, 2000 CLWKR, ACS COW)

plus commute_<year>_repwts.npy for the ACS years, row-aligned with the CSV.

NYC is identified by PUMA code. The 2000-vintage and 2010-vintage codes for the five
boroughs are the same 55 codes (03701-03710 Bronx, 03801-03810 Manhattan, 03901-03903
Staten Island, 04001-04018 Brooklyn, 04101-04114 Queens); the 2000 equivalency file
(census_2000_pumeq5_ny) is read to confirm this rather than assumed. The 1990 PUMAs are
different and are read off the 1990 dictionary's state-dependent codes: 1990 New York
PUMAs for the city are 03701-04114 as well in the 5% file (each borough is a set of
"main PUMA" 3-digit prefixes 037-041); the script verifies by checking that the
place-of-work PUMA codes for Manhattan reconcile with the residence codes.
"""

from __future__ import annotations

import argparse
import re
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd

from _paths import RAW, REFERENCE, INTERIM

OUT_DIR = INTERIM / "acs1"
NYC_COUNTIES = {"005": "Bronx", "047": "Brooklyn", "061": "Manhattan", "081": "Queens", "085": "Staten Island"}

# The 55 PUMA codes of the five boroughs, shared by the 2000 and 2010 vintages.
def _boro_of_puma(code: str) -> str | None:
    c = int(code)
    if 3701 <= c <= 3710: return "Bronx"
    if 3801 <= c <= 3810: return "Manhattan"
    if 3901 <= c <= 3903: return "Staten Island"
    if 4001 <= c <= 4018: return "Brooklyn"
    if 4101 <= c <= 4114: return "Queens"
    return None


# The 1990 5% PUMAs number the boroughs differently. Verified by weighted population
# against the 1990 census counts (Bronx 1.20m, Manhattan 1.49m, Staten Island 0.38m,
# Brooklyn 2.30m, Queens 1.95m): the H-record PUMA prefixes below sum to 1.196m,
# 1.475m, 0.381m, 2.288m and 1.939m, all within 1%.
def _boro_of_puma_1990(code: str) -> str | None:
    p = code[:3]
    return {"050": "Bronx", "051": "Manhattan", "052": "Staten Island", "053": "Brooklyn", "054": "Queens"}.get(p)

# Place-of-work PUMA for Manhattan: 2000 and 2010 vintages use 03800; 2020 uses 04100.
MANHATTAN_POW = {"2000": 3800, "2010": 3800, "2020": 4100}

# JWTR (2005-2018) -> JWTRNS (2019+): streetcar 3->5, subway 4->3, railroad 5->4.
JWTR_TO_JWTRNS = {3: 5, 4: 3, 5: 4}


def race_label(rac1p, hisp) -> np.ndarray:
    r = pd.to_numeric(rac1p, errors="coerce"); h = pd.to_numeric(hisp, errors="coerce")
    lab = np.where(h > 1, "Hispanic", np.where(r == 1, "White", np.where(r == 2, "Black", np.where(r == 6, "Asian", "Other"))))
    return lab


def verify_2000_pumas() -> None:
    """Confirm from the equivalency file that the 55 codes above are the five boroughs."""
    path = REFERENCE / "census_2000_pumeq5_ny.txt"
    if not path.exists():
        print("  note: PUMEQ5 file not on disk; using the documented codes unverified")
        return
    pumas: dict[str, set[str]] = {}
    for line in path.read_text(encoding="latin-1").splitlines():
        # summary level 781: state, superpuma, puma, county
        m = re.match(r"^\s*781\s+(\d{2})\s+(\d{5})\s+(\d{5})\s+(\d{3})", line)
        if m:
            pumas.setdefault(m.group(3), set()).add(m.group(4))
    if not pumas:
        # fall back to the fixed-width form used in some state files
        for line in path.read_text(encoding="latin-1").splitlines():
            if line[:3] == "781":
                pumas.setdefault(line[10:15], set()).add(line[15:18])
    nyc = {p for p, cs in pumas.items() if cs & set(NYC_COUNTIES)}
    ours = {f"{c:05d}" for c in range(3701, 4115) if _boro_of_puma(f"{c:05d}")}
    if nyc and nyc != ours:
        raise SystemExit(f"2000 NYC PUMA set differs from the documented codes: {sorted(nyc ^ ours)[:10]}")
    print(f"  2000 equivalency file confirms {len(nyc) or 'n/a'} NYC PUMAs")


# ------------------------------------------------------------------ ACS 2005-2014


def extract_acs(year: int) -> None:
    zp = RAW / "acs_pums_1yr" / str(year) / "csv_pny.zip"
    if not zp.exists():
        print(f"  {year}: {zp.name} not on disk, skipped"); return
    z = zipfile.ZipFile(zp); member = [n for n in z.namelist() if n.lower().endswith(".csv")][0]
    header = pd.read_csv(z.open(member), nrows=0).columns.tolist()
    up = {h.upper(): h for h in header}
    want = ["SERIALNO", "PUMA", "PWGTP", "AGEP", "SEX", "ESR", "JWTR", "JWMNP", "POWSP", "POWPUMA", "PERNP", "WKHP",
            "RAC1P", "HISP", "SCHL", "OCCP", "NATIVITY", "JWDP", "COW", "INDP", "NAICSP", "WKW", "WKWN", "ADJINC"]
    cols = [up[c] for c in want if c in up]
    reps = [up[f"PWGTP{i}"] for i in range(1, 81) if f"PWGTP{i}" in up]
    frames = []
    for chunk in pd.read_csv(z.open(member), usecols=cols + reps, dtype={up["PUMA"]: str, up.get("OCCP", "OCCP"): str,
                                                                        up.get("POWPUMA", "POWPUMA"): str, up.get("INDP", "INDP"): str, up.get("NAICSP", "NAICSP"): str}, chunksize=200_000, low_memory=False):
        chunk.columns = [c.upper() for c in chunk.columns]
        chunk["PUMA"] = chunk["PUMA"].str.zfill(5)
        boro = chunk["PUMA"].map(_boro_of_puma)
        keep = chunk[boro.notna()].copy(); keep["borough"] = boro[boro.notna()]
        frames.append(keep)
    d = pd.concat(frames, ignore_index=True)
    vintage = "2000" if year <= 2011 else "2010"
    mode = pd.to_numeric(d["JWTR"], errors="coerce").map(lambda v: JWTR_TO_JWTRNS.get(int(v), int(v)) if v == v else np.nan)
    pow_puma = pd.to_numeric(d["POWPUMA"], errors="coerce")
    pow_state = pd.to_numeric(d["POWSP"], errors="coerce")
    # SCHL coding changed in 2008: 1-16 with 13 = BA before, 1-24 with 21 = BA after.
    schl = pd.to_numeric(d["SCHL"], errors="coerce")
    ba = (schl >= 21) if year >= 2008 else (schl >= 13)
    out = pd.DataFrame({
        "year": year, "puma": d["PUMA"], "borough": d["borough"], "weight": d["PWGTP"].astype(int),
        "age": d["AGEP"], "sex": d["SEX"], "esr": pd.to_numeric(d["ESR"], errors="coerce"),
        "mode": mode, "minutes": pd.to_numeric(d["JWMNP"], errors="coerce"),
        "pow_state": pow_state, "pow_puma": pow_puma,
        "pow_manhattan": ((pow_state == 36) & (pow_puma == MANHATTAN_POW[vintage])).astype(int),
        "earnings": pd.to_numeric(d["PERNP"], errors="coerce"), "hours": pd.to_numeric(d["WKHP"], errors="coerce"),
        "race": race_label(d["RAC1P"], d["HISP"]), "foreign_born": (pd.to_numeric(d["NATIVITY"], errors="coerce") == 2).astype(int),
        "ba_plus": ba.astype(int), "occ": d["OCCP"].astype(str), "depart": pd.to_numeric(d["JWDP"], errors="coerce"),
        "ind": d["INDP"].astype(str).str.strip() if "INDP" in d else "", "naics": d["NAICSP"].astype(str).str.strip() if "NAICSP" in d else "",
        "cow": pd.to_numeric(d["COW"], errors="coerce") if "COW" in d else np.nan,
        # Usual weekly hours and whether the person worked 50-52 weeks, both over the past 12
        # months (the earnings reference period). WKWN (weeks, 2019 on) replaced WKW (bands; 1 = 50-52).
        "hours_usual": pd.to_numeric(d["WKHP"], errors="coerce"),
        # The Bureau's within-year income factor (six implied decimals), which puts the rolling
        # twelve-month income of every respondent on the survey year's dollars.
        "adjinc": (pd.to_numeric(d["ADJINC"], errors="coerce") / 1e6) if "ADJINC" in d else np.nan,
        "fullyear": ((pd.to_numeric(d["WKWN"], errors="coerce") >= 50) if "WKWN" in d else (pd.to_numeric(d["WKW"], errors="coerce") == 1) if "WKW" in d else np.nan),
    })
    _write(out, d[[c.upper() for c in reps]].fillna(0).to_numpy(dtype=np.int32) if reps else None, year)


# ------------------------------------------------------------------ Census 2000 5%


def _layout_2000() -> dict[str, dict[str, tuple[int, int]]]:
    """Field name -> (start, width) for the H and P records, from the Bureau's workbook."""
    xls = REFERENCE / "census_2000_pums_record_layout.xls"
    sheets = pd.read_excel(xls, sheet_name=None, header=None)
    out: dict[str, dict[str, tuple[int, int]]] = {"H": {}, "P": {}}
    for name, df in sheets.items():
        rec = "H" if "hous" in name.lower() else "P"
        # Columns: RT, BEG, END, LEN, A/N, VARIABLE, DESCRIPTION, ... (header on row 1).
        for _, row in df.iterrows():
            var = str(row[5]).strip()
            beg = pd.to_numeric(row[1], errors="coerce"); end = pd.to_numeric(row[2], errors="coerce")
            if re.fullmatch(r"[A-Z][A-Z0-9]{1,9}", var) and pd.notna(beg) and pd.notna(end) and end >= beg > 0:
                out[rec].setdefault(var, (int(beg), int(end - beg + 1)))
    return out


def extract_2000() -> None:
    src = RAW / "census_2000_pums" / "all_New_York" / "REVISEDPUMS5_36.TXT"
    if not src.exists():
        alt = list((RAW / "census_2000_pums").rglob("*PUMS5_36.TXT"))
        if not alt:
            print("  2000: PUMS file not on disk, skipped"); return
        src = alt[0]
    lay = _layout_2000()
    P = lay["P"]; H = lay["H"]
    need_p = ["SERIALNO", "PWEIGHT", "AGE", "SEX", "ESR", "TRVMNS", "TRVTIME", "POWST5", "POWPUMA5", "EARNS", "HOURS",
              "RACE1", "HISPAN", "EDUC", "OCCCEN5", "CITIZEN", "INDCEN", "INDNAICS", "CLWKR", "WEEKS"]
    missing = [k for k in need_p if k not in P]
    print(f"  2000 layout: {len(P)} person fields, {len(H)} housing fields; missing {missing}")
    def cut(line, key, table):
        beg, wid = table[key]; return line[beg - 1: beg - 1 + wid]
    rows = []; boro_by_serial: dict[str, tuple[str, str]] = {}
    with open(src, "r", encoding="latin-1") as fh:
        for line in fh:
            if line[0] == "H":
                puma = cut(line, "PUMA5", H); serial = cut(line, "SERIALNO", H)
                b = _boro_of_puma(puma) if puma.strip().isdigit() else None
                if b: boro_by_serial[serial] = (puma, b)
            elif line[0] == "P":
                serial = cut(line, "SERIALNO", P)
                if serial not in boro_by_serial: continue
                puma, b = boro_by_serial[serial]
                rows.append([f for f in (cut(line, k, P) for k in need_p if k in P)] + [puma, b])
    cols = [k for k in need_p if k in P] + ["puma", "borough"]
    d = pd.DataFrame(rows, columns=cols)
    num = lambda c: pd.to_numeric(d[c].str.strip(), errors="coerce")
    esr = num("ESR")
    mode2000 = num("TRVMNS")  # 2000: 1 car, 2 bus, 3 streetcar, 4 subway, 5 railroad, 6 ferry, 7 taxi, 8 motorcycle, 9 bicycle, 10 walked, 11 other, 12 worked at home (blank N/A)
    mode = mode2000.map({1: 1, 2: 2, 3: 5, 4: 3, 5: 4, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 12, 12: 11})
    pow_state = num("POWST5"); pow_puma = num("POWPUMA5")
    race1 = num("RACE1"); hisp = num("HISPAN")
    race = np.where(hisp > 1, "Hispanic", np.where(race1 == 1, "White", np.where(race1 == 2, "Black", np.where(race1 == 4, "Asian", "Other"))))
    out = pd.DataFrame({
        "year": 2000, "puma": d["puma"], "borough": d["borough"], "weight": num("PWEIGHT").fillna(0).astype(int),
        "age": num("AGE"), "sex": num("SEX"), "esr": esr, "mode": mode, "minutes": num("TRVTIME").where(mode.notna() & (mode != 11)),
        "pow_state": pow_state, "pow_puma": pow_puma, "pow_manhattan": ((pow_state == 36) & (pow_puma == MANHATTAN_POW["2000"])).astype(int),
        "earnings": num("EARNS"), "hours": num("HOURS"), "race": race,
        "foreign_born": (num("CITIZEN") >= 4).astype(int) if "CITIZEN" in d else 0,   # 4 naturalized, 5 not a citizen
        "ba_plus": (num("EDUC") >= 13).astype(int), "occ": d["OCCCEN5"].str.strip() if "OCCCEN5" in d else "", "depart": num("TRVDEPT") if "TRVDEPT" in d else np.nan,
        "ind": d["INDCEN"].str.strip() if "INDCEN" in d else "", "naics": d["INDNAICS"].str.strip() if "INDNAICS" in d else "", "cow": num("CLWKR") if "CLWKR" in d else np.nan,
        # The 2000 census asked about 1999: HOURS is usual weekly hours and WEEKS the weeks worked that year.
        "hours_usual": num("HOURS"), "fullyear": (num("WEEKS") >= 50).astype(int),
    })
    _write(out, None, 2000)


# ------------------------------------------------------------------ Census 1990 5%


def _layout_1990() -> dict[str, tuple[int, int]]:
    """Field -> (start, width) from the plain-text dictionary; the P record follows the H record."""
    txt = (REFERENCE / "census_1990_pums_dictionary.txt").read_text(encoding="latin-1")
    out: dict[str, tuple[int, int]] = {}
    rec = "H"
    for line in txt.splitlines():
        m = re.match(r"^D\s+([A-Z0-9]+)\s+(\d+)\s+(\d+)", line)
        if m:
            name, width, beg = m.group(1), int(m.group(2)), int(m.group(3))
            if name == "RECTYPE" and "H:RECTYPE" in out:
                rec = "P"
            out[f"{rec}:{name}"] = (beg, width)
    return out


def extract_1990() -> None:
    files = list((RAW / "census_1990_pums").rglob("*.TXT")) + list((RAW / "census_1990_pums").rglob("*.txt"))
    files = [f for f in files if "PUMSAX" in f.name.upper() or f.stat().st_size > 50e6]
    if not files:
        print("  1990: PUMS file not on disk, skipped"); return
    lay = _layout_1990()
    def get(rec, key): return lay[f"{rec}:{key}"]
    needed = ["SERIALNO", "PWGT1", "AGE", "SEX", "RLABOR", "MEANS", "TRAVTIME", "POWSTATE", "POWPUMA", "INCOME1", "INCOME2", "HOURS",
              "RACE", "HISPANIC", "YEARSCH", "OCCUP", "CITIZEN", "DEPART", "INDUSTRY", "CLASS", "HOUR89", "WEEK89"]
    have = [k for k in needed if f"P:{k}" in lay]
    print(f"  1990 layout: {sum(k.startswith('P:') for k in lay)} person fields; using {have}")
    def cut(line, rec, key):
        beg, wid = get(rec, key); return line[beg - 1: beg - 1 + wid]
    rows = []; cur = None
    with open(files[0], "r", encoding="latin-1") as fh:
        for line in fh:
            if line[0] == "H":
                puma = cut(line, "H", "PUMA"); serial = cut(line, "H", "SERIALNO")
                b = _boro_of_puma_1990(puma)
                cur = (puma, b) if b else None
            elif line[0] == "P" and cur:
                rows.append([cut(line, "P", k) for k in have] + [cur[0], cur[1]])
    d = pd.DataFrame(rows, columns=have + ["puma", "borough"])
    num = lambda c: pd.to_numeric(d[c].str.strip(), errors="coerce") if c in d else pd.Series(np.nan, index=d.index)
    # RLABOR 1990: 1 civilian employed at work, 2 employed not at work, 3 unemployed, 4/5 armed forces, 6 NILF, 0 N/A
    esr = num("RLABOR").replace({0: np.nan})
    means = num("MEANS")  # 1990: 1 car, 2 bus/trolley bus, 3 streetcar, 4 subway, 5 railroad, 6 ferry, 7 taxi, 8 motorcycle, 9 bicycle, 10 walked, 11 other, 12 worked at home
    mode = means.map({1: 1, 2: 2, 3: 5, 4: 3, 5: 4, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 12, 12: 11})
    pow_state = num("POWSTATE"); pow_puma = num("POWPUMA")
    race = num("RACE"); hisp = num("HISPANIC")
    # 1990 RACE: 001 White, 002 Black, 004-... Asian/PI codes 004-036 roughly; Hispanic 001-004 where >0 means Hispanic origin
    race_lab = np.where(hisp > 0, "Hispanic", np.where(race == 1, "White", np.where(race == 2, "Black", np.where((race >= 4) & (race <= 36), "Asian", "Other"))))
    earn = num("INCOME1").fillna(0) + num("INCOME2").fillna(0)
    out = pd.DataFrame({
        "year": 1990, "puma": d["puma"], "borough": d["borough"], "weight": num("PWGT1").fillna(0).astype(int),
        "age": num("AGE"), "sex": num("SEX"), "esr": esr, "mode": mode, "minutes": num("TRAVTIME").where(mode.notna() & (mode != 11)),
        "pow_state": pow_state, "pow_puma": pow_puma, "pow_manhattan": 0,
        "earnings": earn, "hours": num("HOURS"), "race": race_lab,
        "foreign_born": (num("CITIZEN").isin([1, 2, 3])).astype(int),   # 0 born in US; 1 born in PR/outlying; 2 born abroad of US parents; 3 naturalized; 4 not a citizen
        "ba_plus": (num("YEARSCH") >= 14).astype(int), "occ": d["OCCUP"].str.strip() if "OCCUP" in d else "", "depart": num("DEPART"),
        "ind": d["INDUSTRY"].str.strip() if "INDUSTRY" in d else "", "naics": "", "cow": num("CLASS") if "CLASS" in d else np.nan,
        # `hours` above is hours worked LAST WEEK (the commute questions' reference week). The
        # earnings year is 1989, so the like-for-like measures are HOUR89 (usual weekly hours
        # in 1989) and WEEK89 (weeks worked in 1989), matching the later censuses and the ACS.
        "hours_usual": num("HOUR89").replace({0: np.nan}), "fullyear": (num("WEEK89") >= 50).astype(int),
    })
    out["foreign_born"] = (num("CITIZEN") >= 3).astype(int)
    # Manhattan place of work in 1990: the modal POWPUMA among workers living in Manhattan who work in-state.
    man = out[(out.borough == "Manhattan") & (out.pow_state == 36)]
    if len(man):
        code = int(man.groupby("pow_puma").weight.sum().idxmax())
        out["pow_manhattan"] = ((out.pow_state == 36) & (out.pow_puma == code)).astype(int)
        print(f"  1990 Manhattan place-of-work code inferred as {code}")
    _write(out, None, 1990)


def _write(out: pd.DataFrame, reps, year: int) -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    out.to_csv(OUT_DIR / f"commute_{year}.csv.gz", index=False, compression="gzip")
    if reps is not None:
        np.save(OUT_DIR / f"commute_{year}_repwts.npy", reps)
    w = out[(out.esr == 1) & out.minutes.notna()]
    med = np.nan
    if len(w):
        o = np.argsort(w.minutes.to_numpy()); cw = np.cumsum(w.weight.to_numpy()[o]); med = w.minutes.to_numpy()[o][np.searchsorted(cw, cw[-1] / 2)]
    print(f"  {year}: {len(out):,} NYC person rows, {out.weight.sum():,} weighted; commuters {w.weight.sum():,}, median minutes {med}")


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--years", nargs="*", type=int, default=None)
    args = ap.parse_args(argv)
    years = args.years or [1990, 2000] + list(range(2005, 2015))
    if any(2005 <= y <= 2011 for y in years):
        verify_2000_pumas()
    for y in years:
        if y == 1990: extract_1990()
        elif y == 2000: extract_2000()
        else: extract_acs(y)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
