#!/usr/bin/env python3
"""Build the earnings-accountability threshold table.

Public Law 119-21 sec. 84001 (HEA 454(c)) and the Department of Education's
final rule of July 1, 2026 (91 FR 40136, FR Doc. 2026-13286) compare a
program's completer earnings with the median earnings of working adults aged
25 to 34 in the state where the institution is located (or nationally). The
rule's regulatory impact analysis names six thresholds and the method:

  In-state HS, National HS, Same-state same-field BA, Same-state BA,
  National same-field BA, National BA.
  Working = positive income from wages, salary or self-employment (incl. farm).
  Not currently enrolled in college. Field = 2-digit CIP. Weighted medians.
  ACS 5-year PUMS, earnings in the file's dollar year (ADJINC). The rule sets
  the same-state same-field BA threshold to $1 where fewer than 16 individuals
  are available (34 CFR 668.2, "earnings threshold", paragraph (3)).

Two sources are combined:
  1. The College Scorecard institution file publishes EARN_THR_STATE and
     EARN_THR_NAT (the undergraduate thresholds, 2024 ACS 5-year since the
     October 15, 2026 release; since June 10, 2026 the derivation excludes
     the unemployed and those not in the labor force).
  2. ACS 2024 5-year PUMS via the Census API, to replicate the HS thresholds
     under both working definitions and to compute the four bachelor's
     thresholds the Scorecard does not publish.

Field of study follows the Department's own crosswalk: Appendix Table 1 of the PPD:2026 technical
appendix maps the ACS general degree field (the first two digits of FOD1P) to a 2-digit CIP code.
The detailed crosswalk used in the first build is kept as a sensitivity column.

--year 2023 --inflate 1.029494 --tag _acs2023 rebuilds the benchmarks on the vintage the Department
used for PPD:2026 (2023 ACS 5-year, CPI-U to 2024 dollars); the *_employed columns then follow the
PPD:2026 working definition (in the labor force, employed, positive earnings), for reconciliation.

Usage:
  CENSUS_API_KEY=... python3 build_earnings_thresholds.py \
      --institution-csv _private/data/scorecard/Most-Recent-Cohorts-Institution_11022026.csv \
      --outdir data/earnings-accountability --date 20261105
"""
import argparse, csv, json, os, sys, time, urllib.request, urllib.parse
from collections import defaultdict

YEAR = "2024"
def api(pr=False): return f"https://api.census.gov/data/{YEAR}/acs/acs5/pums" + ("pr" if pr else "")
INFLATE = 1.0  # CPI-U factor applied to every median (1.0 when the file year is the dollar year)
# PPD:2026 technical appendix, Appendix Table 1: ACS general degree field -> 2-digit CIP
DEPT_XW = {11:"01",13:"03",14:"04",15:"05",19:"09",20:"10",21:"11",22:"12",23:"13",24:"14",25:"15",26:"16",29:"19",32:"22",33:"23",34:"24",35:"25",36:"26",37:"27",38:"29",40:"30",41:"31",48:"38",49:"39",50:"40",51:"41",52:"42",53:"43",54:"44",55:"45",56:"46",57:"47",59:"49",60:"50",61:"51",62:"52",64:"54"}
VARS = ["PWGTP", "WAGP", "SEMP", "ADJINC", "SCH", "FOD1P", "ESR", "SCHL", "AGEP"]
CIP2_TITLES = {
 "01":"Agricultural and related sciences","03":"Natural resources and conservation","04":"Architecture and related services",
 "05":"Area, ethnic, cultural, gender and group studies","09":"Communication, journalism and related programs","10":"Communications technologies",
 "11":"Computer and information sciences","12":"Culinary, entertainment and personal services","13":"Education","14":"Engineering",
 "15":"Engineering and engineering-related technologies","16":"Foreign languages, literatures and linguistics","19":"Family and consumer sciences",
 "22":"Legal professions and studies","23":"English language and literature","24":"Liberal arts and sciences, general studies and humanities",
 "25":"Library science","26":"Biological and biomedical sciences","27":"Mathematics and statistics","29":"Military technologies and applied sciences",
 "30":"Multi/interdisciplinary studies","31":"Parks, recreation, leisure, fitness and kinesiology","38":"Philosophy and religious studies",
 "39":"Theology and religious vocations","40":"Physical sciences","41":"Science technologies","42":"Psychology",
 "43":"Homeland security, law enforcement, firefighting and protective services","44":"Public administration and social service professions",
 "45":"Social sciences","46":"Construction trades","47":"Mechanic and repair technologies","48":"Precision production",
 "49":"Transportation and materials moving","50":"Visual and performing arts","51":"Health professions and related programs",
 "52":"Business, management, marketing and related support services","54":"History"}
TERRITORIES = ["AS", "GU", "MP", "VI", "FM", "MH", "PW"]
FIPS_TO_ABBR = {"01":"AL","02":"AK","04":"AZ","05":"AR","06":"CA","08":"CO","09":"CT","10":"DE","11":"DC","12":"FL","13":"GA","15":"HI","16":"ID","17":"IL","18":"IN","19":"IA","20":"KS","21":"KY","22":"LA","23":"ME","24":"MD","25":"MA","26":"MI","27":"MN","28":"MS","29":"MO","30":"MT","31":"NE","32":"NV","33":"NH","34":"NJ","35":"NM","36":"NY","37":"NC","38":"ND","39":"OH","40":"OK","41":"OR","42":"PA","44":"RI","45":"SC","46":"SD","47":"TN","48":"TX","49":"UT","50":"VT","51":"VA","53":"WA","54":"WV","55":"WI","56":"WY","72":"PR"}

def fetch(params, key, tries=4, base=None):
    q = dict(params); q["key"] = key
    url = base + "?" + urllib.parse.urlencode(q, safe=":,")
    for t in range(tries):
        try:
            with urllib.request.urlopen(url, timeout=300) as r:
                if r.status == 204: return []
                data = json.loads(r.read().decode("utf-8"))
                hdr = data[0]; return [dict(zip(hdr, row)) for row in data[1:]]
        except Exception as e:
            if t == tries - 1: raise
            time.sleep(5 * (t + 1))

def pull(state_fips, schl, key):
    rows = fetch({"get": ",".join(VARS), "for": f"state:{state_fips}", "AGEP": "25:34", "SCHL": schl}, key,
                 base=api(pr=(state_fips == "72")))
    out = []
    for r in rows:
        adj = float(r["ADJINC"])  # the API serves ADJINC as a decimal factor (e.g. 1.193241)
        wag = float(r["WAGP"]) if r["WAGP"] not in ("", "N") else 0.0
        sem = float(r["SEMP"]) if r["SEMP"] not in ("", "N") else 0.0
        earn = (wag + sem) * adj
        out.append({"w": int(r["PWGTP"]), "earn": earn, "sch": r["SCH"], "esr": r["ESR"], "fod": r["FOD1P"], "schl": r["SCHL"]})
    return out

def r_half_up(v):
    """Round half up; Python's round() rounds half to even, which is off by one dollar on .5 values."""
    return int(v + 0.5)

def to_dollars(med):
    """Round the median to whole dollars (halves up), then apply the dollar factor and round again.
    With INFLATE = 1 (the 2020-24 file, already in 2024 dollars) this is a single rounding."""
    m = r_half_up(med)
    return m if INFLATE == 1.0 else r_half_up(m * INFLATE)

def wmedian(rows):
    """Weighted median of earn by w. Returns (median, n_unweighted, n_weighted).
    Sort by earnings and accumulate PWGTP; the median is the first record whose running weight
    reaches half the total. If the running weight equals exactly half, the median is the average
    of that record and the next one (the standard weighted median; PEER Center uses the same rule)."""
    if not rows: return (None, 0, 0)
    rows = sorted(rows, key=lambda x: x["earn"])
    tot = sum(x["w"] for x in rows); half = tot / 2.0; cum = 0
    for i, x in enumerate(rows):
        cum += x["w"]
        if cum == half and i + 1 < len(rows): return (to_dollars((x["earn"] + rows[i + 1]["earn"]) / 2.0), len(rows), tot)
        if cum >= half: return (to_dollars(x["earn"]), len(rows), tot)
    return (to_dollars(rows[-1]["earn"]), len(rows), tot)

def working_A(x):  # rule RIA footnote 58: positive earnings, enrollment excluded
    return x["earn"] > 0 and x["sch"] == "1"
def working_B(x):  # Scorecard June 2026 derivation: in the labor force and employed, positive earnings, not enrolled
    return x["esr"] in ("1", "2", "4", "5") and x["earn"] > 0 and x["sch"] == "1"

def load_crosswalk(path):
    xw = {}
    with open(path, encoding="utf-8") as f:
        for r in csv.DictReader(f):
            xw[str(int(r["fod"]))] = str(int(float(r["cip4"]))).zfill(2)
    return xw

def load_scorecard(path):
    thr = {}
    with open(path, encoding="utf-8") as f:
        for r in csv.DictReader(f):
            st = r["STABBR"]; v = r.get("EARN_THR_STATE", ""); n = r.get("EARN_THR_NAT", "")
            if st and v not in ("", "NULL", "NA", "PrivacySuppressed"): thr.setdefault(st, set()).add(int(float(v)))
            if n not in ("", "NULL", "NA", "PrivacySuppressed"): thr.setdefault("US", set()).add(int(float(n)))
    return {k: (sorted(v)[0] if len(v) == 1 else None) for k, v in thr.items()}

def main():
    global YEAR, INFLATE
    ap = argparse.ArgumentParser()
    ap.add_argument("--institution-csv", required=True)
    ap.add_argument("--crosswalk", default="data/ai-exposure/acs_fod_to_cip4_crosswalk_20261028.csv")
    ap.add_argument("--outdir", default="data/earnings-accountability")
    ap.add_argument("--date", default=time.strftime("%Y%m%d"))
    ap.add_argument("--year", default="2024")
    ap.add_argument("--inflate", type=float, default=1.0)
    ap.add_argument("--tag", default="")
    ap.add_argument("--states", default="")  # comma list of FIPS for testing
    a = ap.parse_args()
    YEAR, INFLATE = a.year, a.inflate
    key = os.environ.get("CENSUS_API_KEY")
    if not key: sys.exit("CENSUS_API_KEY not set")
    os.makedirs(a.outdir, exist_ok=True)
    xw = load_crosswalk(a.crosswalk)
    sc = load_scorecard(a.institution_csv) if a.year == "2024" else {}
    fips = list(FIPS_TO_ABBR)  # 50 states, DC, and PR (PR from the PRCS endpoint; excluded from the national pool)
    if a.states: fips = a.states.split(",")
    basis_file = f"{a.year} 5-year PUMS"
    hs_all, ba_all = [], []
    state_rows, cellsA, cellsB, cellsD = [], {}, {}, {}
    natA, natB = defaultdict(list), defaultdict(list)
    natA_pr, ba_pr = defaultdict(list), []  # diagnostic only: national bachelor's pools with Puerto Rico added
    t0 = time.time()
    for i, f in enumerate(fips, 1):
        ab = FIPS_TO_ABBR[f]
        hs = pull(f, "16,17", key); ba = pull(f, "21", key)
        if f != "72": hs_all += hs; ba_all += ba
        ba_pr += ba
        hsA = wmedian([x for x in hs if working_A(x)]); hsB = wmedian([x for x in hs if working_B(x)])
        baA = wmedian([x for x in ba if working_A(x)]); baB = wmedian([x for x in ba if working_B(x)])
        state_rows.append({"state": ab, "state_fips": f,
            "hs_threshold_scorecard": sc.get(ab, ""), "hs_threshold_national_scorecard": sc.get("US", ""),
            "hs_median_acs_positive_earnings": hsA[0], "hs_n_positive_earnings": hsA[1],
            "hs_median_acs_employed": hsB[0], "hs_n_employed": hsB[1],
            "ba_median_acs_positive_earnings": baA[0], "ba_n_positive_earnings": baA[1],
            "ba_median_acs_employed": baB[0], "ba_n_employed": baB[1],
            "basis": ("PRCS " if f == "72" else "ACS ") + basis_file,
            "hs_matches_scorecard": ("not published" if sc.get(ab) is None else ("yes" if sc.get(ab) == hsA[0] else "no"))})
        gA, gB, gD = defaultdict(list), defaultdict(list), defaultdict(list)
        for x in ba:
            try: gen = int(x["fod"]) // 100
            except ValueError: continue
            c2 = DEPT_XW.get(gen); cd = xw.get(str(int(x["fod"])))
            if working_A(x):
                if c2:
                    gA[c2].append(x)
                    if f != "72": natA[c2].append(x)
                    natA_pr[c2].append(x)
                if cd: gD[cd].append(x)
            if working_B(x) and c2:
                gB[c2].append(x)
                if f != "72": natB[c2].append(x)
        for c2 in gA: cellsA[(ab, f, c2)] = wmedian(gA[c2])
        for c2 in gB: cellsB[(ab, c2)] = wmedian(gB[c2])
        for c2 in gD: cellsD[(ab, c2)] = wmedian(gD[c2])
        print(f"[{i}/{len(fips)}] {ab} hs n={len(hs)} ba n={len(ba)} hsA={hsA[0]} scorecard={sc.get(ab)} ba={baA[0]} {time.time()-t0:.0f}s", flush=True)
    nhA = wmedian([x for x in hs_all if working_A(x)]); nhB = wmedian([x for x in hs_all if working_B(x)])
    nbA = wmedian([x for x in ba_all if working_A(x)]); nbB = wmedian([x for x in ba_all if working_B(x)])
    state_rows.append({"state": "US", "state_fips": "", "hs_threshold_scorecard": sc.get("US", ""), "hs_threshold_national_scorecard": sc.get("US", ""),
        "hs_median_acs_positive_earnings": nhA[0], "hs_n_positive_earnings": nhA[1], "hs_median_acs_employed": nhB[0], "hs_n_employed": nhB[1],
        "ba_median_acs_positive_earnings": nbA[0], "ba_n_positive_earnings": nbA[1], "ba_median_acs_employed": nbB[0], "ba_n_employed": nbB[1],
        "basis": f"ACS {basis_file}, 50 states and DC", "hs_matches_scorecard": ("yes" if sc.get("US") == nhA[0] else "no") if sc else "not published"})
    # Territories and Freely Associated States: no ACS or PRCS data. 34 CFR 668.2 "earnings threshold" paragraph (3)
    # sets the threshold to one dollar where the Census data are not available (91 FR 40281; discussion at 91 FR 40174).
    if not a.states:
        for ab in TERRITORIES:
            state_rows.append({"state": ab, "state_fips": "", "hs_threshold_scorecard": "", "hs_threshold_national_scorecard": sc.get("US", ""),
                "hs_median_acs_positive_earnings": 1, "hs_n_positive_earnings": 0, "hs_median_acs_employed": "", "hs_n_employed": 0,
                "ba_median_acs_positive_earnings": 1, "ba_n_positive_earnings": 0, "ba_median_acs_employed": "", "ba_n_employed": 0,
                "basis": "No Census data; one dollar under 34 CFR 668.2 earnings threshold paragraph (3)", "hs_matches_scorecard": "not published"})
    nfA = {c2: wmedian(rows) for c2, rows in natA.items()}; nfB = {c2: wmedian(rows) for c2, rows in natB.items()}
    ba_state = {r["state"]: (r["ba_median_acs_positive_earnings"], r["ba_median_acs_employed"]) for r in state_rows}
    field_rows = []
    for (ab, f, c2), (m, n, wn) in sorted(cellsA.items(), key=lambda kv: (fips.index(kv[0][1]), kv[0][2])):
        mB, nB, _ = cellsB.get((ab, c2), (None, 0, 0)); mD, nD, _ = cellsD.get((ab, c2), (None, 0, 0))
        nf = nfA.get(c2, (None, 0, 0)); nfb = nfB.get(c2, (None, 0, 0))
        rule = m if n >= 16 else 1
        cands = [v for v in (ba_state[ab][0], rule, nf[0]) if v is not None]
        field_rows.append({"state": ab, "state_fips": f, "cip2": c2, "cip2_title": CIP2_TITLES.get(c2, ""),
            "ba_same_state_same_field_median": m, "n_unweighted": n, "n_weighted": wn, "threshold_rule_value": rule,
            "below_16_respondents": "yes" if n < 16 else "no", "below_30_respondents": "yes" if n < 30 else "no",
            "ba_national_same_field_median": nf[0], "ba_national_same_field_n": nf[1], "ba_same_state_median": ba_state[ab][0],
            "graduate_threshold_lowest_of_three": min(cands) if cands else None,
            "ba_same_state_same_field_median_employed": mB, "n_unweighted_employed": nB,
            "ba_national_same_field_median_employed": nfb[0], "ba_same_state_median_employed": ba_state[ab][1],
            "ba_same_state_same_field_median_detailed_crosswalk": mD, "n_unweighted_detailed_crosswalk": nD})
    for c2, (m, n, wn) in sorted(nfA.items()):
        nfb = nfB.get(c2, (None, 0, 0))
        field_rows.append({"state": "US", "state_fips": "", "cip2": c2, "cip2_title": CIP2_TITLES.get(c2, ""),
            "ba_same_state_same_field_median": m, "n_unweighted": n, "n_weighted": wn, "threshold_rule_value": m,
            "below_16_respondents": "yes" if n < 16 else "no", "below_30_respondents": "yes" if n < 30 else "no",
            "ba_national_same_field_median": m, "ba_national_same_field_n": n, "ba_same_state_median": nbA[0],
            "graduate_threshold_lowest_of_three": min(v for v in (nbA[0], m) if v is not None),
            "ba_same_state_same_field_median_employed": nfb[0], "n_unweighted_employed": nfb[1],
            "ba_national_same_field_median_employed": nfb[0], "ba_same_state_median_employed": nbB[0],
            "ba_same_state_same_field_median_detailed_crosswalk": "", "n_unweighted_detailed_crosswalk": ""})
    p1 = os.path.join(a.outdir, f"earnings_thresholds_by_state_{a.date}{a.tag}.csv")
    p2 = os.path.join(a.outdir, f"earnings_thresholds_ba_by_state_field_{a.date}{a.tag}.csv")
    with open(p1, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(state_rows[0].keys())); w.writeheader(); w.writerows(state_rows)
    with open(p2, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(field_rows[0].keys())); w.writeheader(); w.writerows(field_rows)
    p3 = os.path.join(a.outdir, f"earnings_thresholds_ba_national_with_pr_diagnostic_{a.date}{a.tag}.csv")
    with open(p3, "w", newline="") as f:
        w = csv.writer(f); w.writerow(["cip2", "cip2_title", "ba_national_median_excluding_pr", "ba_national_median_including_pr"])
        w.writerow(["all", "All fields", nbA[0], wmedian([x for x in ba_pr if working_A(x)])[0]])
        for c2 in sorted(nfA): w.writerow([c2, CIP2_TITLES.get(c2, ""), nfA[c2][0], wmedian(natA_pr[c2])[0]])
    print("wrote", p1, len(state_rows), "rows;", p2, len(field_rows), "rows;", p3, flush=True)

if __name__ == "__main__":
    main()