#!/usr/bin/env python3
"""F1 stage 1b: cells with borrower state, for the geography control. Stdlib only.

Same universe and field rules as build_cube.py. Aggregates
  cells_state.csv.gz   draw, band, bgroup, month, state -> loans, chargedoff   (all loans)
  lender_state.csv.gz  orig_lender, draw, band, bgroup, month, state -> loans, chargedoff
                       (only lenders named in <lender_list>, one name per line)
Usage: build_state_cells.py <ppp_data_dir> <lender_list> <out_dir>
"""
import csv, glob, gzip, os, sys, time
from collections import defaultdict
from build_cube import cents, band_of, SOLE


def main(src, lender_list, out):
    keep = set(x.rstrip("\n") for x in open(lender_list, encoding="utf-8") if x.strip())
    cells = defaultdict(lambda: [0, 0]); lc = defaultdict(lambda: [0, 0]); t0 = time.time()
    for f in sorted(glob.glob(os.path.join(src, "*.csv"))):
        with open(f, encoding="utf-8", errors="replace", newline="") as fh:
            r = csv.reader(fh)
            h = [x.lstrip("") for x in next(r)]
            ix = {k: i for i, k in enumerate(h)}
            iD, iSt, iPm, iCur, iUnd = ix["DateApproved"], ix["LoanStatus"], ix["ProcessingMethod"], ix["CurrentApprovalAmount"], ix["UndisbursedAmount"]
            iBt, iOl, iBs = ix["BusinessType"], ix["OriginatingLender"], ix["BorrowerState"]
            for row in r:
                cur = cents(row[iCur])
                if cents(row[iUnd]) >= cur:
                    continue
                d = row[iD]
                pm = row[iPm].strip()
                draw = "1" if pm == "PPP" else ("2" if pm == "PPS" else "?")
                bt = row[iBt].strip()
                bg = "SOLE" if bt in SOLE else ("BLANK" if bt == "" else "OTHER")
                k = (draw, band_of(cur), bg, d[6:10] + "-" + d[0:2], row[iBs].strip())
                co = 1 if row[iSt].strip() == "Charged Off" else 0
                c = cells[k]; c[0] += 1; c[1] += co
                ol = row[iOl].strip()
                if ol in keep:
                    x = lc[(ol,) + k]; x[0] += 1; x[1] += co
        print(f"{os.path.basename(f)} t={time.time()-t0:.0f}s", flush=True)
    with gzip.open(os.path.join(out, "cells_state.csv.gz"), "wt", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh); w.writerow(["draw", "band", "bgroup", "month", "state", "loans", "chargedoff"])
        for k in sorted(cells):
            w.writerow(list(k) + cells[k])
    with gzip.open(os.path.join(out, "lender_state.csv.gz"), "wt", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh); w.writerow(["orig_lender", "draw", "band", "bgroup", "month", "state", "loans", "chargedoff"])
        for k in sorted(lc):
            w.writerow(list(k) + lc[k])
    open(os.path.join(out, "STAGE1B.DONE"), "w").write(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n")


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