#!/usr/bin/env python3
"""Screen College Scorecard field-of-study earnings against the earnings-accountability thresholds.

One snapshot, not the official test. The official measure (34 CFR 668.402, 668.403; 91 FR 40136)
uses the fourth tax year after completion for the 2020-21 cohort first, pools small cohorts, and
needs two failing years out of three. This screen compares the Scorecard's published four-year
median earnings (2017-18 and 2018-19 completers, measured in calendar 2022 and 2023, 2024 dollars)
with the threshold each credential
level is assigned in Table 5.5 of the rule's impact analysis (91 FR 40213):

  undergraduate certificate, associate, bachelor's, post-baccalaureate certificate:
      in-state high-school median, or the national high-school median
  master's, doctoral, first professional, graduate certificate:
      lowest of same-state same-field BA, same-state BA, national same-field BA;
      or lowest of national same-field BA and national BA

Which column applies depends on the institution's share of in-state students (FAFSA addresses),
which the Department published in PPD:2026 (outstate_flag); both columns are kept for every program.

Vintage. 34 CFR 668.2 sets the threshold from Census data "during the year of the associated measured
earnings." The Scorecard earnings were measured in 2022 and 2023, so the primary thresholds are the
2019-2023 ACS 5-year file in 2024 dollars (--thresholds, --fields). The 2020-2024 file, which the
Scorecard has displayed since October 2026, is kept as a sensitivity column (--alt-thresholds, --alt-fields).

Exposure columns. t4_completers_2024_25 is Title IV completers in award year 2024-25 (PPD:2026,
t4completers_2025). loan_borrowers and loan_debt_total are Stafford and Grad PLUS borrowers among the
2022-23 and 2023-24 completers and their total debt (Scorecard DEBT_ALL_STGP_ANY_N times _MEAN).

Usage:
  python3 build_earnings_accountability_screen.py --date 20261106
"""
import argparse, csv, os
from collections import defaultdict

UNDERGRAD = {"1", "2", "3", "4"}
GRAD = {"5", "6", "7", "8"}
# Table 5.22, 91 FR 40236: twenty 6-digit CIP codes with delayed implementation. The Scorecard
# file is 4-digit, so the flag is by family: "all" where every listed code fills the family,
# "part" where the family also holds programs outside the list.
TIPPED = {"1204": "most", "1206": "all", "5135": "all", "1205": "part", "5209": "part"}

PPD_LEV = {"Undergrad Certificate": "1", "Associate": "2", "Bachelor": "3", "Post-Bacc Certificate": "4", "Master's": "5", "Doctoral": "6", "First Professional Degree": "7", "Graduate Certificate": "8"}

def num(v):
    return None if v in ("", "NULL", "NA", "PS", "PrivacySuppressed") else float(v)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--fos", default="_private/data/scorecard/Most-Recent-Cohorts-Field-of-Study_10152026.csv")
    ap.add_argument("--inst", default="_private/data/scorecard/Most-Recent-Cohorts-Institution_11022026.csv")
    ap.add_argument("--thresholds", default="data/earnings-accountability/earnings_thresholds_by_state_20261106_acs2023.csv")
    ap.add_argument("--fields", default="data/earnings-accountability/earnings_thresholds_ba_by_state_field_20261106_acs2023.csv")
    ap.add_argument("--alt-thresholds", default="data/earnings-accountability/earnings_thresholds_by_state_20261106.csv")
    ap.add_argument("--alt-fields", default="data/earnings-accountability/earnings_thresholds_ba_by_state_field_20261106.csv")
    ap.add_argument("--outdir", default="data/earnings-accountability")
    ap.add_argument("--ppd-inst", default="_private/data/ed-ppd2026/ppd2026_institution.csv")
    ap.add_argument("--date", required=True)
    a = ap.parse_args()

    # Department's out-of-state flag (PPD:2026, 2017-18 and 2018-19 Title IV enrollees); state column where absent
    oflag = {}; ppd_state = {}; dfail = {}
    if os.path.exists(a.ppd_inst):
        for r in csv.DictReader(open(a.ppd_inst)):
            oflag[r["opeid6"].zfill(6)] = r["outstate_flag"]; ppd_state[r["opeid6"].zfill(6)] = r["stabbr"]
    ppd_test = a.ppd_inst.replace("ppd2026_institution.csv", "ppd2026_earnings_test.csv")
    if os.path.exists(ppd_test):
        for r in csv.DictReader(open(ppd_test)):
            lv = PPD_LEV.get(r["credlev"].strip())
            try: k = (r["opeid6"].zfill(6), str(int(float(r["cip4"]))).zfill(4), lv)
            except ValueError: continue
            if lv is None: continue
            v = r["mstr_obbb_fail_cip2_wageb"]
            dfail[k] = "" if v == "" else ("yes" if float(v) == 1 else "no")
    t4c = {}
    if os.path.exists(a.ppd_inst):
        for r in csv.DictReader(open(a.ppd_inst)):
            lv = PPD_LEV.get(r["credlev"].strip())
            try: k = (r["opeid6"].zfill(6), str(int(float(r["cip4"]))).zfill(4), lv)
            except ValueError: continue
            v = r["t4completers_2025"]
            try: t4c[k] = int(float(v)) if lv else None
            except ValueError: pass  # blank, NA, or "privacy suppressed"
            if t4c.get(k) is None: t4c.pop(k, None)
    st_of = {}
    with open(a.inst, encoding="utf-8") as f:
        for r in csv.DictReader(f): st_of[r["UNITID"]] = r["STABBR"]
    def make(tpath, fpath):
        T = {r["state"]: r for r in csv.DictReader(open(tpath))}
        hs_nat = int(T["US"]["hs_median_acs_positive_earnings"]); ba_nat = int(T["US"]["ba_median_acs_positive_earnings"])
        cell = {}; natf = {}
        for r in csv.DictReader(open(fpath)):
            if r["state"] == "US": natf[r["cip2"]] = int(r["ba_same_state_same_field_median"])
            else: cell[(r["state"], r["cip2"])] = r

        def thresholds(state, credlev, cip2):
            """Return (type, in_state_value, national_value, note)."""
            t = T.get(state)
            if credlev in UNDERGRAD:
                ins = int(t["hs_median_acs_positive_earnings"]) if t else None
                return ("high school", ins, hs_nat, "" if t and ins != 1 else ("no Census data, one dollar" if ins == 1 else "state not in table"))
            ba_state = int(t["ba_median_acs_positive_earnings"]) if t else None
            c = cell.get((state, cip2)); nf = natf.get(cip2)
            note = ""
            if ba_state == 1: sf = 1; note = "no Census data, one dollar"
            elif nf is None:
                # no ACS degree field maps to this CIP family (Appendix Table 1): there is no same-field
                # comparison, so the same-state bachelor's median applies (as PEER's "other" field rows)
                sf = None; note = "field outside the Department's crosswalk; same-state bachelor's median"
            elif c is None: sf = 1; note = "no ACS respondents in the state-field cell, one dollar"
            elif c["below_16_respondents"] == "yes": sf = 1; note = "fewer than 16 respondents in the state-field cell, one dollar"
            else: sf = int(c["ba_same_state_same_field_median"])
            cands = [v for v in (ba_state, sf, nf) if v is not None]
            ins = min(cands) if cands else None
            nat = min(v for v in (nf, ba_nat) if v is not None)
            return ("bachelor's, lowest of three", ins, nat, note)
        return T, thresholds

    T, thresholds = make(a.thresholds, a.fields)
    T_alt, thresholds_alt = make(a.alt_thresholds, a.alt_fields)

    seen = set(); rows = []
    summ = defaultdict(lambda: defaultdict(int))
    with open(a.fos, encoding="utf-8") as f:
        for r in csv.DictReader(f):
            if r["MAIN"] != "1" or r["CONTROL"] == "Foreign": continue
            key = (r["OPEID6"], r["CIPCODE"], r["CREDLEV"])
            if key in seen: continue
            seen.add(key)
            state = st_of.get(r["UNITID"]) or ppd_state.get(r["OPEID6"].zfill(6), "")
            if not state or state not in T:
                summ[("", r["CREDLEV"], r["CREDDESC"], r["CONTROL"])]["programs"] += 1
                continue
            cip4 = r["CIPCODE"].zfill(4); cip2 = cip4[:2]
            earn = num(r["EARN_MDN_4YR"]); wne = num(r["EARN_COUNT_WNE_4YR"])
            k = (state, r["CREDLEV"], r["CREDDESC"], r["CONTROL"])
            summ[k]["programs"] += 1
            if earn is None: continue
            ttype, ins, nat, note = thresholds(state, r["CREDLEV"], cip2)
            b_ins = (earn < ins) if ins is not None else None
            b_nat = earn < nat
            of = oflag.get(r["OPEID6"].zfill(6), "")
            use_nat = of in ("1", "1.0")
            thr_app = nat if use_nat else ins
            b_app = b_nat if use_nat else b_ins
            _, ins2, nat2, _ = thresholds_alt(state, r["CREDLEV"], cip2)
            thr_alt = nat2 if use_nat else ins2
            b_alt = (earn < thr_alt) if thr_alt is not None else None
            pk = (r["OPEID6"].zfill(6), cip4, r["CREDLEV"])
            comp = t4c.get(pk)
            ln = num(r["DEBT_ALL_STGP_ANY_N"]); lm = num(r["DEBT_ALL_STGP_ANY_MEAN"])
            debt = int(round(ln * lm)) if ln is not None and lm is not None else None
            tip = TIPPED.get(cip4, "no")
            summ[k]["with_earnings"] += 1
            summ[k]["cohort_wne"] += int(wne or 0)
            if b_ins: summ[k]["below_in_state"] += 1; summ[k]["cohort_wne_below_in_state"] += int(wne or 0)
            if b_nat: summ[k]["below_national"] += 1
            if b_app:
                summ[k]["below_applied"] += 1; summ[k]["cohort_wne_below_applied"] += int(wne or 0)
                summ[k]["t4_completers_2024_25_below_applied"] += comp or 0; summ[k]["loan_debt_below_applied"] += debt or 0
            if b_alt: summ[k]["below_applied_acs2024"] += 1
            summ[k]["t4_completers_2024_25"] += comp or 0
            if b_ins and tip in ("all", "most"): summ[k]["below_in_state_tipped_delay"] += 1
            if ins == 1: summ[k]["one_dollar_threshold"] += 1
            rows.append({"unitid": r["UNITID"], "opeid6": r["OPEID6"], "institution": r["INSTNM"], "state": state, "control": r["CONTROL"],
                "cip4": cip4, "cip_title": r["CIPDESC"].rstrip("."), "credlev": r["CREDLEV"], "credential": r["CREDDESC"],
                "earn_mdn_4yr": int(earn), "cohort_working_not_enrolled": "" if wne is None else int(wne),
                "threshold_type": ttype, "threshold_in_state": ins, "threshold_national": nat,
                "below_in_state": "" if b_ins is None else ("yes" if b_ins else "no"), "below_national": "yes" if b_nat else "no",
                "gap_in_state": "" if ins is None else int(earn - ins), "tipped_delay_family": tip, "threshold_note": note,
                "dept_outstate_flag": "" if of == "" else int(float(of)), "threshold_applied": thr_app,
                "below_applied": "" if b_app is None else ("yes" if b_app else "no"), "gap_applied": "" if thr_app is None else int(earn - thr_app),
                "dept_mstr_fail": dfail.get(pk, ""),
                "threshold_applied_acs2024": "" if thr_alt is None else thr_alt, "below_applied_acs2024": "" if b_alt is None else ("yes" if b_alt else "no"),
                "t4_completers_2024_25": "" if comp is None else comp, "loan_borrowers_2022_24": "" if ln is None else int(ln),
                "loan_debt_total_2022_24": "" if debt is None else debt})
            if dfail.get((r["OPEID6"].zfill(6), cip4, r["CREDLEV"])) == "yes": summ[k]["dept_fail"] += 1
    p1 = os.path.join(a.outdir, f"program_screen_{a.date}.csv")
    with open(p1, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys())); w.writeheader(); w.writerows(rows)
    cols = ["programs", "with_earnings", "below_in_state", "below_national", "below_applied", "cohort_wne_below_applied", "dept_fail", "below_applied_acs2024", "t4_completers_2024_25", "t4_completers_2024_25_below_applied", "loan_debt_below_applied", "below_in_state_tipped_delay", "one_dollar_threshold", "cohort_wne", "cohort_wne_below_in_state"]
    p2 = os.path.join(a.outdir, f"state_screen_summary_{a.date}.csv")
    with open(p2, "w", newline="") as f:
        w = csv.writer(f); w.writerow(["state", "credlev", "credential", "control"] + cols)
        for k in sorted(summ): w.writerow(list(k) + [summ[k][c] for c in cols])
    print("wrote", p1, len(rows), "rows;", p2, len(summ), "rows")

if __name__ == "__main__":
    main()