#!/usr/bin/env python3
"""Independent second read of F1 (pandas C parser, vectorised; shares no code with reproduce_chargeoffs.py).

Recomputes the headline numbers straight from the 13 raw SBA files and compares them with
f1/out/numbers.json. Prints one line per check and a final PASS/FAIL count.
Usage: check_chargeoffs.py <ppp_data_dir> <numbers.json>
"""
import glob, json, os, sys
import numpy as np
import pandas as pd

src, numbers = sys.argv[1], sys.argv[2]
N = json.load(open(numbers))
COLS = ["DateApproved", "LoanStatus", "ProcessingMethod", "CurrentApprovalAmount", "UndisbursedAmount",
        "BusinessType", "OriginatingLender", "ServicingLenderName", "BorrowerState"]
EDGES = [-np.inf, 500000, 1000000, 2083333, 5000000, 15000000, 35000000, 199999999, np.inf]
parts1, parts2, parts3 = [], [], []
rows_read = 0
for f in sorted(glob.glob(os.path.join(src, "*.csv"))):
    for ch in pd.read_csv(f, usecols=COLS, dtype=str, keep_default_na=False, encoding="utf-8",
                          encoding_errors="replace", chunksize=750000):
        rows_read += len(ch)
        cur = (pd.to_numeric(ch["CurrentApprovalAmount"].str.strip().replace("", "0")) * 100).round().astype("int64")
        und = (pd.to_numeric(ch["UndisbursedAmount"].str.strip().replace("", "0")) * 100).round().astype("int64")
        ch = ch[und < cur].copy(); cur = cur[und < cur]
        ch["cents"] = cur
        ch["lender"] = ch["OriginatingLender"].str.strip()
        ch["serv"] = ch["ServicingLenderName"].str.strip()
        ch["st"] = ch["LoanStatus"].str.strip()
        ch["co"] = (ch["st"] == "Charged Off").astype("int64")
        ch["e4"] = (ch["st"] == "Exemption 4").astype("int64")
        ch["draw"] = ch["ProcessingMethod"].str.strip().map({"PPP": 1, "PPS": 2}).fillna(0).astype("int64")
        ch["band"] = pd.cut(ch["cents"], EDGES, labels=False, right=True) + 1
        bt = ch["BusinessType"].str.strip()
        ch["bg"] = np.where(bt.isin(["Sole Proprietorship", "Self-Employed Individuals", "Independent Contractors"]), "S",
                            np.where(bt == "", "B", "O"))
        ch["month"] = ch["DateApproved"].str.slice(6, 10) + "-" + ch["DateApproved"].str.slice(0, 2)
        ch["state"] = ch["BorrowerState"].str.strip()
        ch["co_cents"] = ch["cents"] * ch["co"]
        ch["one"] = 1
        parts1.append(ch.groupby(["lender", "draw", "band", "bg", "month", "state"], sort=False)[["one", "co", "e4", "cents", "co_cents"]].sum())
        parts3.append(ch.groupby(["serv"], sort=False)[["one", "co"]].sum())
        parts2.append(ch.groupby(["lender", "serv"], sort=False)[["one"]].sum())
g = pd.concat(parts1).groupby(level=[0, 1, 2, 3, 4, 5]).sum().reset_index()
sv = pd.concat(parts3).groupby(level=0).sum()
pairs = pd.concat(parts2).groupby(level=[0, 1]).sum().reset_index()
res = []


def chk(name, got, want, tol=1e-9):
    ok = (got == want) if isinstance(want, (int, str)) else (want is not None and abs(got - want) <= tol * max(1, abs(want)))
    res.append(ok)
    print(("PASS " if ok else "FAIL ") + name + f": check={got!r} reproduce={want!r}")


def q(v):
    v = np.sort(np.asarray(v, dtype=float))
    return dict(min=v[0], median=float(np.median(v)), max=v[-1], p10=float(np.percentile(v, 10)), p90=float(np.percentile(v, 90)))


chk("rows_read", int(rows_read), N["rows_read"])
chk("universe_loans", int(g["one"].sum()), N["universe_loans"])
chk("charged_off_loans", int(g["co"].sum()), N["charged_off_loans"])
chk("charged_off_usd", float(g["co_cents"].sum()) / 100, N["charged_off_usd"])
L = g.groupby("lender")[["one", "co", "e4", "cents", "co_cents"]].sum()
L["rate"] = L["co"] / L["one"]
chk("lender_keys", int(len(L)), N["lender_keys"])
big = L[L["one"] >= 20000]
chk("ge20k lenders", int(len(big)), N["ge20k"]["lenders"])
chk("ge5k lenders", int((L["one"] >= 5000).sum()), N["ge5k"]["lenders"])
chk("ge20k loans", int(big["one"].sum()), N["ge20k"]["loans"])
s = q(big["rate"])
for k in ("min", "median", "max", "p10", "p90"):
    chk("ge20k rate " + k, s[k], N["ge20k"]["rate"][k])
chk("ge20k max lender", str(big["rate"].idxmax()), N["ge20k"]["rate"]["max_lender"])
chk("ge20k min lender", str(big["rate"].idxmin()), N["ge20k"]["rate"]["min_lender"])
s = q(big["co_cents"] / big["cents"])
chk("ge20k usd rate max", s["max"], N["ge20k"]["usd_rate"]["max"])
top = L.sort_values(["co"], ascending=False).head(10)
chk("top10 share of charge-offs", top["co"].sum() / L["co"].sum(), N["top10_by_charge_offs"]["share_of_charge_offs"])
chk("top10 share of loans", top["one"].sum() / L["one"].sum(), N["top10_by_charge_offs"]["share_of_all_loans"])
chk("top10 share of co usd", top["co_cents"].sum() / L["co_cents"].sum(), N["top10_by_charge_offs"]["share_of_charged_off_usd"])
D = g.groupby("draw")[["one", "co"]].sum()
for d in (1, 2):
    chk(f"draw{d} rate", D.loc[d, "co"] / D.loc[d, "one"], N["draw"][str(d)]["rate"])
bl = set(big.index)
gb = g[g["lender"].isin(bl)]
for d in (1, 2):
    x = gb[gb["draw"] == d].groupby("lender")[["one", "co"]].sum(); x = x[x["one"] >= 1000]
    s = q(x["co"] / x["one"])
    chk(f"ge20k draw{d} min", s["min"], N["ge20k"][f"draw{d}"]["min"]); chk(f"ge20k draw{d} max", s["max"], N["ge20k"][f"draw{d}"]["max"])
for nm, mask in (("small_loans", gb["band"] <= 3), ("sole_prop_ic", gb["bg"] == "S"), ("other_borrowers", gb["bg"] == "O")):
    x = gb[mask].groupby("lender")[["one", "co"]].sum(); x = x[x["one"] >= 1000]
    s = q(x["co"] / x["one"])
    for k in ("min", "median", "max"):
        chk(f"ge20k {nm} {k}", s[k], N["ge20k"][nm][k])
# mix adjustment without state: all-lender cell rates, and leave-one-out
c = g.groupby(["lender", "draw", "band", "bg", "month"])[["one", "co"]].sum().reset_index()
cell = c.groupby(["draw", "band", "bg", "month"])[["one", "co"]].sum().rename(columns={"one": "cn", "co": "cc"}).reset_index()
c = c.merge(cell, on=["draw", "band", "bg", "month"])
c["exp_all"] = c["one"] * c["cc"] / c["cn"]
loo_n, loo_c = c["cn"] - c["one"], c["cc"] - c["co"]
c["exp_loo"] = np.where(loo_n > 0, c["one"] * loo_c / loo_n.where(loo_n > 0, 1), 0.0)
E = c[c["lender"].isin(bl)].groupby("lender")[["co", "exp_all", "exp_loo"]].sum()
s = q(E["co"] / E["exp_all"])
for k in ("min", "median", "max", "p10", "p90"):
    chk("ge20k O/E " + k, s[k], N["ge20k"]["mix_adjusted_oe"][k], 1e-6)
if "mix_adjusted_oe_loo" in N["ge20k"]:
    s = q(E["co"] / E["exp_loo"])
    for k in ("min", "median", "max", "p10", "p90"):
        chk("ge20k O/E LOO " + k, s[k], N["ge20k"]["mix_adjusted_oe_loo"][k], 1e-6)
if "mix_state_oe_loo" in N["ge20k"]:
    cs = g.groupby(["draw", "band", "bg", "month", "state"])[["one", "co"]].sum().rename(columns={"one": "cn", "co": "cc"}).reset_index()
    x = g[g["lender"].isin(bl)].merge(cs, on=["draw", "band", "bg", "month", "state"])
    ln, lc = x["cn"] - x["one"], x["cc"] - x["co"]
    x["exp"] = np.where(ln > 0, x["one"] * lc / ln.where(ln > 0, 1), 0.0)
    Es = x.groupby("lender")[["co", "exp"]].sum()
    s = q(Es["co"] / Es["exp"])
    for k in ("min", "median", "max", "p10", "p90"):
        chk("ge20k O/E state LOO " + k, s[k], N["ge20k"]["mix_state_oe_loo"][k], 1e-6)
s = q(big["e4"] / big["one"])
chk("ge20k exemption4 share max", s["max"], N["ge20k"]["exemption4_share"]["max"])
M = g.groupby("month")[["one", "co"]].sum()
bad = [m for m in M.index if abs(M.loc[m, "co"] / M.loc[m, "one"] - N["by_month"][m]["rate"]) > 1e-12 or int(M.loc[m, "one"]) != N["by_month"][m]["loans"]]
chk("by_month rows agree", len(bad), 0)
svb = sv[sv["one"] >= 20000]
s = q(svb["co"] / svb["one"])
chk("servicers ge20k n", int(len(svb)), N["attribution"]["servicers_ge20k"]["n_lenders"])
chk("servicers ge20k max", s["max"], N["attribution"]["servicers_ge20k"]["max"])
chk("loans serviced by other named lender", int(pairs.loc[pairs["lender"] != pairs["serv"], "one"].sum()), N["attribution"]["loans_serviced_by_other_named_lender"])
print(f"RESULT {sum(res)} PASS / {len(res) - sum(res)} FAIL")