#!/usr/bin/env python3
"""Reproduce every F1 number: PPP charge-off rate by originating lender (SBA PPP FOIA release 2024-09-30).

Stage 1 (raw -> cube) is build_cube.py; it runs here when the cube is missing or --rebuild is given.
Stage 2 (cube -> tables + numbers.json) is below. Stdlib only.

Definitions
- Universe: all loans in the 13 FOIA CSVs except the 4 whose UndisbursedAmount >= CurrentApprovalAmount.
- Charged off: LoanStatus == "Charged Off" exactly as SBA records it in this release.
- Lender: OriginatingLender, the name SBA records for the lender that made the loan (not the servicer).
- Draw: ProcessingMethod PPP = first draw, PPS = second draw.
- Size bands: RP3 bands on CurrentApprovalAmount; "small" = bands 1-3 (up to $20,833.33).
- Borrower group: SOLE = Sole Proprietorship, Self-Employed Individuals, Independent Contractors.
- Mix-adjusted expectation: for each lender, sum over cells (draw x band x borrower group x approval month)
  of its loans times the all-lender charge-off rate in that cell; O/E = actual / expected.
Usage: reproduce_chargeoffs.py <ppp_data_dir> <work_f1_dir> [--rebuild]
"""
import csv, gzip, json, os, statistics, sys
from collections import defaultdict

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
BAND_LABEL = {1: "$1-$5,000", 2: "$5,000.01-$10,000", 3: "$10,000.01-$20,833.33", 4: "$20,833.34-$50,000",
              5: "$50,000.01-$150,000", 6: "$150,000.01-$350,000", 7: "$350,000.01-$1,999,999.99", 8: "$2,000,000 and above"}
MIN_SUB = 1000  # a lender enters a sub-population comparison only with >= 1,000 loans in it


def load_cube(d):
    with gzip.open(os.path.join(d, "cube.csv.gz"), "rt", encoding="utf-8", newline="") as fh:
        r = csv.reader(fh); next(r)
        for L, draw, band, bg, month, st, n, c, fg in r:
            yield L, draw, int(band), bg, month, st, int(n), int(c), int(fg)


def rate(co, n):
    return co / n if n else None


def spread(vals):
    """vals: list of (rate, name). Returns summary dict."""
    v = sorted(vals)
    rs = [x[0] for x in v]
    q = statistics.quantiles(rs, n=10, method="inclusive") if len(rs) >= 2 else [None] * 9
    return {"n_lenders": len(rs), "min": rs[0], "min_lender": v[0][1], "median": statistics.median(rs),
            "max": rs[-1], "max_lender": v[-1][1], "p10": q[0], "p90": q[8],
            "max_over_min": (rs[-1] / rs[0]) if rs[0] else None, "p90_over_p10": (q[8] / q[0]) if q[0] else None}


def main(src, wd, rebuild=False):
    cube_dir = os.path.join(wd, "cube")
    if rebuild or not os.path.exists(os.path.join(cube_dir, "STAGE1.DONE")):
        import build_cube
        build_cube.main(src, cube_dir)
    meta = json.load(open(os.path.join(cube_dir, "meta.json")))
    out = os.path.join(wd, "out"); os.makedirs(out, exist_ok=True)

    tot = [0, 0, 0, 0]  # loans, co, cents, co_cents
    status = defaultdict(int)
    L = defaultdict(lambda: defaultdict(int))
    cell = defaultdict(lambda: [0, 0])          # (draw, band, bg, month) -> loans, co
    Lcell = defaultdict(lambda: defaultdict(lambda: [0, 0]))  # lender -> cell -> loans, co
    month = defaultdict(lambda: [0, 0, 0, 0])   # (month, draw) -> loans, co, cents, co_cents
    bandbg = defaultdict(lambda: [0, 0])        # (band, bg, draw) -> loans, co
    for Ln, draw, band, bg, mo, st, n, c, fg in load_cube(cube_dir):
        co = n if st == "C" else 0
        tot[0] += n; tot[1] += co; tot[2] += c; tot[3] += c if st == "C" else 0
        status[st] += n
        x = L[Ln]
        x["loans"] += n; x["co"] += co; x["cents"] += c; x["co_cents"] += c if st == "C" else 0
        x["e4"] += n if st == "E" else 0; x["paid"] += n if st == "P" else 0
        x["d%s_loans" % draw] += n; x["d%s_co" % draw] += co
        if band <= 3:
            x["small_loans"] += n; x["small_co"] += co
        if bg == "SOLE":
            x["sole_loans"] += n; x["sole_co"] += co
        elif bg == "OTHER":
            x["other_loans"] += n; x["other_co"] += co
        k = (draw, band, bg, mo)
        cell[k][0] += n; cell[k][1] += co
        Lcell[Ln][k][0] += n; Lcell[Ln][k][1] += co
        m = month[(mo, draw)]; m[0] += n; m[1] += co; m[2] += c; m[3] += c if st == "C" else 0
        b = bandbg[(band, bg, draw)]; b[0] += n; b[1] += co
    crate = {k: v[1] / v[0] for k, v in cell.items()}
    for Ln, cells in Lcell.items():
        L[Ln]["expected_co"] = sum(v[0] * crate[k] for k, v in cells.items())
        # leave-one-out: the cell rate among all OTHER lenders' loans
        L[Ln]["expected_co_loo"] = sum(v[0] * (cell[k][1] - v[1]) / (cell[k][0] - v[0]) for k, v in cells.items() if cell[k][0] > v[0])
    # geography control (stage 1b): cells add BorrowerState; leave-one-out; lenders with >= 5,000 loans only
    if rebuild or not os.path.exists(os.path.join(cube_dir, "STAGE1B.DONE")):
        import build_state_cells
        with open(os.path.join(cube_dir, "lenders_ge5k.txt"), "w", encoding="utf-8") as fh:
            fh.write("".join(k + "\n" for k, x in sorted(L.items()) if x["loans"] >= 5000))
        build_state_cells.main(src, os.path.join(cube_dir, "lenders_ge5k.txt"), cube_dir)
    cs = {}
    with gzip.open(os.path.join(cube_dir, "cells_state.csv.gz"), "rt", encoding="utf-8", newline="") as fh:
        r = csv.reader(fh); next(r)
        for dr, b, bg, mo, stt, n, co in r:
            cs[(dr, b, bg, mo, stt)] = (int(n), int(co))
    with gzip.open(os.path.join(cube_dir, "lender_state.csv.gz"), "rt", encoding="utf-8", newline="") as fh:
        r = csv.reader(fh); next(r)
        for Ln, dr, b, bg, mo, stt, n, co in r:
            n, co = int(n), int(co); tn, tc = cs[(dr, b, bg, mo, stt)]
            if tn > n:
                L[Ln]["expected_co_state"] += n * (tc - co) / (tn - n)
            L[Ln]["state_loans"] += n
    # servicing-lender attribution
    serv = defaultdict(lambda: [0, 0]); moved = defaultdict(int); moved_total = 0
    with gzip.open(os.path.join(cube_dir, "pairs.csv.gz"), "rt", encoding="utf-8", newline="") as fh:
        r = csv.reader(fh); next(r)
        for ol, sl, n, co in r:
            n, co = int(n), int(co)
            serv[sl][0] += n; serv[sl][1] += co
            if sl != ol:
                moved[ol] += n; moved_total += n

    # ---- lender table (all lenders)
    rows = []
    for Ln, x in L.items():
        e = x["expected_co"]
        rows.append({
            "originating_lender": Ln, "loans": x["loans"], "charged_off_loans": x["co"],
            "charge_off_rate": round(x["co"] / x["loans"], 6),
            "approved_usd": round(x["cents"] / 100, 2), "charged_off_usd": round(x["co_cents"] / 100, 2),
            "charged_off_usd_share": round(x["co_cents"] / x["cents"], 6) if x["cents"] else "",
            "paid_in_full_loans": x["paid"], "exemption4_loans": x["e4"],
            "first_draw_loans": x["d1_loans"], "first_draw_charged_off": x["d1_co"],
            "second_draw_loans": x["d2_loans"], "second_draw_charged_off": x["d2_co"],
            "small_loans_to_20833": x["small_loans"], "small_loans_charged_off": x["small_co"],
            "sole_prop_ic_loans": x["sole_loans"], "sole_prop_ic_charged_off": x["sole_co"],
            "other_borrower_loans": x["other_loans"], "other_borrower_charged_off": x["other_co"],
            "expected_charged_off_mix_adjusted": round(e, 2),
            "actual_to_expected": round(x["co"] / e, 4) if e else "",
            "expected_charged_off_other_lenders": round(x["expected_co_loo"], 2),
            "actual_to_expected_other_lenders": round(x["co"] / x["expected_co_loo"], 4) if x["expected_co_loo"] else "",
            "expected_charged_off_other_lenders_same_state": round(x["expected_co_state"], 2) if x["state_loans"] else "",
            "actual_to_expected_other_lenders_same_state": round(x["co"] / x["expected_co_state"], 4) if x["state_loans"] and x["expected_co_state"] else "",
            "loans_serviced_by_other_named_lender": moved.get(Ln, 0),
            "ge_20000_loans": int(x["loans"] >= 20000), "ge_5000_loans": int(x["loans"] >= 5000),
            "_rate": x["co"] / x["loans"], "_usd": (x["co_cents"] / x["cents"]) if x["cents"] else None,
            "_oe": (x["co"] / e) if e else None, "_oe_loo": (x["co"] / x["expected_co_loo"]) if x["expected_co_loo"] else None,
            "_oe_state": (x["co"] / x["expected_co_state"]) if x["state_loans"] and x["expected_co_state"] else None,
        })
    rows.sort(key=lambda r: (-r["loans"], r["originating_lender"]))
    for i, r in enumerate(rows, 1):
        r["rank_by_loans"] = i
    cols = ["rank_by_loans"] + [k for k in rows[0] if k != "rank_by_loans" and not k.startswith("_")]
    with open(os.path.join(out, "ppp-charge-offs-by-lender.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.DictWriter(fh, fieldnames=cols, extrasaction="ignore"); w.writeheader(); w.writerows(rows)

    # ---- month table
    with open(os.path.join(out, "ppp-charge-offs-by-approval-month.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["approval_month", "draw", "loans", "charged_off_loans", "charge_off_rate", "approved_usd", "charged_off_usd"])
        for (mo, dr) in sorted(month):
            v = month[(mo, dr)]
            w.writerow([mo, {"1": "first", "2": "second"}.get(dr, dr), v[0], v[1], round(v[1] / v[0], 6), round(v[2] / 100, 2), round(v[3] / 100, 2)])
    # ---- band x borrower group x draw table
    with open(os.path.join(out, "ppp-charge-offs-by-size-and-borrower.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["band", "band_label", "borrower_group", "draw", "loans", "charged_off_loans", "charge_off_rate"])
        for (band, bg, dr) in sorted(bandbg):
            v = bandbg[(band, bg, dr)]
            w.writerow([band, BAND_LABEL[band], bg, {"1": "first", "2": "second"}.get(dr, dr), v[0], v[1], round(v[1] / v[0], 6)])

    # ---- numbers for the page
    N = {"source": "SBA PPP FOIA loan-level release 2024-09-30", "rows_read": meta["rows_read"],
         "excluded_rows": meta["excluded_rows"], "universe_loans": tot[0], "charged_off_loans": tot[1],
         "charge_off_rate": tot[1] / tot[0], "approved_usd": tot[2] / 100, "charged_off_usd": tot[3] / 100,
         "status_counts": dict(status), "max_status_date": meta["max_status_date"],
         "approval_dates": [meta["min_date_approved"], meta["max_date_approved"]], "lender_keys": len(L)}
    big = [r for r in rows if r["loans"] >= 20000]
    mid = [r for r in rows if r["loans"] >= 5000]
    for tag, pop in (("ge20k", big), ("ge5k", mid)):
        N[tag] = {"lenders": len(pop), "loans": sum(r["loans"] for r in pop),
                  "charged_off": sum(r["charged_off_loans"] for r in pop),
                  "rate": spread([(r["_rate"], r["originating_lender"]) for r in pop]),
                  "usd_rate": spread([(r["_usd"], r["originating_lender"]) for r in pop])}
    N["ge20k"]["share_of_universe_loans"] = N["ge20k"]["loans"] / tot[0]
    N["ge20k"]["share_of_all_charge_offs"] = N["ge20k"]["charged_off"] / tot[1]
    # concentration: top 10 lenders by charged-off loans (all lenders)
    top = sorted(rows, key=lambda r: (-r["charged_off_loans"], r["originating_lender"]))[:10]
    N["top10_by_charge_offs"] = {
        "lenders": [r["originating_lender"] for r in top],
        "share_of_charge_offs": sum(r["charged_off_loans"] for r in top) / tot[1],
        "share_of_charged_off_usd": sum(r["charged_off_usd"] for r in top) / (tot[3] / 100),
        "share_of_all_loans": sum(r["loans"] for r in top) / tot[0]}
    # draw split, universe
    d = {"1": [0, 0], "2": [0, 0]}
    for (mo, dr), v in month.items():
        d[dr][0] += v[0]; d[dr][1] += v[1]
    N["draw"] = {k: {"loans": v[0], "charged_off": v[1], "rate": v[1] / v[0]} for k, v in d.items()}
    # draw split within the >=20k lenders
    for dr in ("1", "2"):
        vals = [(r["d_co"] / r["d_n"], r["originating_lender"]) for r in
                ({"d_co": r["%s_draw_charged_off" % ("first" if dr == "1" else "second")],
                  "d_n": r["%s_draw_loans" % ("first" if dr == "1" else "second")],
                  "originating_lender": r["originating_lender"]} for r in big) if r["d_n"] >= MIN_SUB]
        N["ge20k"]["draw%s" % dr] = spread(vals)
    # controls within the >=20k lenders
    for key, n_col, c_col in (("small_loans", "small_loans_to_20833", "small_loans_charged_off"),
                              ("sole_prop_ic", "sole_prop_ic_loans", "sole_prop_ic_charged_off"),
                              ("other_borrowers", "other_borrower_loans", "other_borrower_charged_off")):
        vals = [(r[c_col] / r[n_col], r["originating_lender"]) for r in big if r[n_col] >= MIN_SUB]
        N["ge20k"][key] = spread(vals)
    oe = [(r["_oe"], r["originating_lender"]) for r in big]
    N["ge20k"]["mix_adjusted_oe"] = spread(oe)
    N["ge20k"]["oe_ge_2"] = sum(1 for v, _ in oe if v >= 2)
    for tag, col in (("mix_adjusted_oe_loo", "_oe_loo"), ("mix_state_oe_loo", "_oe_state")):
        v = [(r[col], r["originating_lender"]) for r in big]
        N["ge20k"][tag] = spread(v)
        N["ge20k"][tag]["n_ge_2"] = sum(1 for a, _ in v if a >= 2)
        N["ge20k"][tag]["n_le_half"] = sum(1 for a, _ in v if a <= 0.5)
        N["ge20k"][tag]["top5"] = [[n, a] for a, n in sorted(v, reverse=True)[:5]]
        N["ge20k"][tag]["bottom5"] = [[n, a] for a, n in sorted(v)[:5]]
    # sensitivity: merge spelling variants of the same name (case, punctuation, '&' vs 'and')
    import re
    def nv(s):
        return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9 ]", " ", s.lower().replace("&", " and "))).strip()
    mg = defaultdict(lambda: [0, 0, []])
    for r in rows:
        m = mg[nv(r["originating_lender"])]; m[0] += r["loans"]; m[1] += r["charged_off_loans"]; m[2].append(r["originating_lender"])
    mb = [(v[1] / v[0], " / ".join(v[2])) for v in mg.values() if v[0] >= 20000]
    N["ge20k"]["name_variants_merged"] = spread(mb)
    N["ge20k"]["name_variants_merged"]["merged_groups_ge20k"] = [n for _, n in mb if " / " in n]
    N["ge20k"]["oe_le_half"] = sum(1 for v, _ in oe if v <= 0.5)
    # rank correlation raw rate vs O/E (Spearman, average ranks)
    def ranks(xs):
        o = sorted(range(len(xs)), key=lambda i: xs[i]); rk = [0.0] * len(xs); i = 0
        while i < len(o):
            j = i
            while j + 1 < len(o) and xs[o[j + 1]] == xs[o[i]]:
                j += 1
            for k in range(i, j + 1):
                rk[o[k]] = (i + j) / 2 + 1
            i = j + 1
        return rk
    a = ranks([r["_rate"] for r in big]); b = ranks([r["_oe"] for r in big])
    ma, mb = statistics.mean(a), statistics.mean(b)
    N["ge20k"]["spearman_rate_vs_oe"] = sum((x - ma) * (y - mb) for x, y in zip(a, b)) / (
        (sum((x - ma) ** 2 for x in a) * sum((y - mb) ** 2 for y in b)) ** 0.5)
    e4 = [(r["exemption4_loans"] / r["loans"], r["originating_lender"]) for r in big]
    N["ge20k"]["exemption4_share"] = spread(e4)
    N["ge20k"]["charge_off_plus_e4_rate"] = spread(
        [((r["charged_off_loans"] + r["exemption4_loans"]) / r["loans"], r["originating_lender"]) for r in big])
    # attribution
    N["attribution"] = {"loans_serviced_by_other_named_lender": moved_total,
                        "share": moved_total / tot[0],
                        "servicers_ge20k": spread([(v[1] / v[0], k) for k, v in serv.items() if v[0] >= 20000])}
    N["ge20k"]["serviced_elsewhere_share"] = spread(
        [(r["loans_serviced_by_other_named_lender"] / r["loans"], r["originating_lender"]) for r in big])
    # month summary
    mon = defaultdict(lambda: [0, 0])
    for (mo, dr), v in month.items():
        mon[mo][0] += v[0]; mon[mo][1] += v[1]
    N["by_month"] = {mo: {"loans": v[0], "charged_off": v[1], "rate": v[1] / v[0]} for mo, v in sorted(mon.items())}
    json.dump(N, open(os.path.join(out, "numbers.json"), "w"), indent=1, sort_keys=True)
    print(json.dumps({k: N[k] for k in ("universe_loans", "charged_off_loans", "charge_off_rate", "charged_off_usd")}))


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], "--rebuild" in sys.argv)