#!/usr/bin/env python3
"""F1 stage 1: SBA PPP FOIA loan-level release (2024-09-30) -> compact aggregate cube. Stdlib only.

Universe (house filter, same as RP3-SMALL-LOANS-0919): every row of the 13 CSVs except rows whose
UndisbursedAmount >= CurrentApprovalAmount (in integer cents). Read with csv + errors="replace"
(a duckdb read silently drops 11 rows when OriginatingLender is projected; RP3 measured this).

Outputs in <out_dir>:
  cube.csv.gz    orig_lender, draw, band, bgroup, month, status -> loans, cents, forgiven_cents
  pairs.csv.gz   orig_lender, serv_lender -> loans, chargedoff_loans
  statusdate.csv status, status_month -> loans
  meta.json      per-file row counts, excluded rows and cents, header check, value vocabularies
  STAGE1.DONE    marker
Usage: build_cube.py <ppp_data_dir> <out_dir>
"""
import csv, glob, gzip, json, os, sys, time
from collections import defaultdict, Counter

SOLE = {"Sole Proprietorship", "Self-Employed Individuals", "Independent Contractors"}
# RP3 bands, upper bounds inclusive, integer cents; band 8 = $2,000,000 and above
BAND_UPPER = [500000, 1000000, 2083333, 5000000, 15000000, 35000000, 199999999]
STATUS = {"Paid in Full": "P", "Charged Off": "C", "Exemption 4": "E"}


def cents(s):
    s = (s or "").strip()
    return int(round(float(s) * 100)) if s else 0


def band_of(c):
    for i, u in enumerate(BAND_UPPER):
        if c <= u:
            return i + 1
    return 8


def main(src, out):
    os.makedirs(out, exist_ok=True)
    t0 = time.time()
    cube = defaultdict(lambda: [0, 0, 0])
    pairs = defaultdict(lambda: [0, 0])
    sdate = Counter()
    meta = {"files": {}, "excluded_rows": 0, "excluded_cents": 0, "rows_read": 0,
            "status_values": Counter(), "draw_values": Counter(), "btype_values": Counter(),
            "header_identical": True, "max_status_date": "", "min_date_approved": "", "max_date_approved": ""}
    header0 = None
    min_d, max_d, max_sd = "9999-99-99", "0000-00-00", "0000-00-00"
    files = sorted(glob.glob(os.path.join(src, "*.csv")))
    for f in files:
        n = 0
        with open(f, encoding="utf-8", errors="replace", newline="") as fh:
            r = csv.reader(fh)
            h = [x.lstrip("") for x in next(r)]
            if header0 is None:
                header0 = h
            elif h != header0:
                meta["header_identical"] = False
            ix = {k: i for i, k in enumerate(h)}
            iD, iSt, iSd, iPm, iCur, iUnd = ix["DateApproved"], ix["LoanStatus"], ix["LoanStatusDate"], ix["ProcessingMethod"], ix["CurrentApprovalAmount"], ix["UndisbursedAmount"]
            iBt, iOl, iSl, iFg = ix["BusinessType"], ix["OriginatingLender"], ix["ServicingLenderName"], ix["ForgivenessAmount"]
            for row in r:
                n += 1
                cur = cents(row[iCur])
                und = cents(row[iUnd])
                if und >= cur:
                    meta["excluded_rows"] += 1
                    meta["excluded_cents"] += cur
                    continue
                d = row[iD]  # MM/DD/YYYY
                iso = d[6:10] + "-" + d[0:2] + "-" + d[3:5]
                if iso < min_d: min_d = iso
                if iso > max_d: max_d = iso
                st_raw = row[iSt].strip()
                st = STATUS.get(st_raw, "O")
                meta["status_values"][st_raw] += 1
                pm = row[iPm].strip()
                meta["draw_values"][pm] += 1
                draw = "1" if pm == "PPP" else ("2" if pm == "PPS" else "?")
                bt = row[iBt].strip()
                meta["btype_values"][bt] += 1
                bg = "SOLE" if bt in SOLE else ("BLANK" if bt == "" else "OTHER")
                ol = row[iOl].strip()
                sl = row[iSl].strip()
                fg = cents(row[iFg])
                k = (ol, draw, band_of(cur), bg, iso[:7], st)
                c = cube[k]
                c[0] += 1; c[1] += cur; c[2] += fg
                p = pairs[(ol, sl)]
                p[0] += 1
                if st == "C":
                    p[1] += 1
                sd = row[iSd].strip()
                sdi = (sd[6:10] + "-" + sd[0:2]) if len(sd) >= 10 else ""
                sdate[(st, sdi)] += 1
                if sd and len(sd) >= 10:
                    sdf = sd[6:10] + "-" + sd[0:2] + "-" + sd[3:5]
                    if sdf > max_sd: max_sd = sdf
        meta["files"][os.path.basename(f)] = n
        meta["rows_read"] += n
        print(f"{os.path.basename(f)} rows={n} t={time.time()-t0:.0f}s", flush=True)
    meta["min_date_approved"], meta["max_date_approved"], meta["max_status_date"] = min_d, max_d, max_sd
    for k in ("status_values", "draw_values", "btype_values"):
        meta[k] = dict(meta[k].most_common())
    meta["cube_cells"] = len(cube)
    meta["pairs"] = len(pairs)
    meta["seconds"] = round(time.time() - t0)
    with gzip.open(os.path.join(out, "cube.csv.gz"), "wt", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["orig_lender", "draw", "band", "bgroup", "month", "status", "loans", "cents", "forgiven_cents"])
        for k in sorted(cube):
            w.writerow(list(k) + cube[k])
    with gzip.open(os.path.join(out, "pairs.csv.gz"), "wt", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["orig_lender", "serv_lender", "loans", "chargedoff_loans"])
        for k in sorted(pairs):
            w.writerow(list(k) + pairs[k])
    with open(os.path.join(out, "statusdate.csv"), "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["status", "status_month", "loans"])
        for k in sorted(sdate):
            w.writerow(list(k) + [sdate[k]])
    json.dump(meta, open(os.path.join(out, "meta.json"), "w"), indent=1)
    open(os.path.join(out, "STAGE1.DONE"), "w").write(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")


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