"""
Who pays late? — a repayment-propensity model for chase prioritisation.

Data: UCI "Default of Credit Card Clients" (Yeh & Lien, 2009) — 30,000 accounts, Taiwan, 2005.
Target: default payment next month (1 = missed).

Pipeline
  1. load + clean (undocumented category codes folded)
  2. features from six months of payment history (no target leakage: everything is known before the month scored)
  3. stratified 80/20 split; 5-fold CV on the training half
  4. logistic baseline vs gradient boosting; calibration; threshold economics
  5. fairness check on protected attributes (excluded from the model, reported by group)
  6. export: trees → JSON for in-browser scoring, evaluation arrays → JSON for the case page
"""
import json, math, sys
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_predict
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.metrics import roc_auc_score, average_precision_score, brier_score_loss, roc_curve, precision_recall_curve
from sklearn.calibration import calibration_curve

SEED = 42
rng = np.random.default_rng(SEED)

# ---------------------------------------------------------------- 1. load + clean
df = pd.read_csv("data/uci_credit_default.csv")
df = df.rename(columns={"default payment next month": "target", "PAY_0": "PAY_1"})
# undocumented codes: EDUCATION 0/5/6 → "other" (4); MARRIAGE 0 → "other" (3)
df["EDUCATION"] = df["EDUCATION"].replace({0: 4, 5: 4, 6: 4})
df["MARRIAGE"] = df["MARRIAGE"].replace({0: 3})
y = df["target"].astype(int).values

PAY = [f"PAY_{i}" for i in range(1, 7)]            # PAY_1 = most recent month (Sept 2005) ... PAY_6 = April
BILL = [f"BILL_AMT{i}" for i in range(1, 7)]
PAMT = [f"PAY_AMT{i}" for i in range(1, 7)]

# ---------------------------------------------------------------- 2. features (all pre-outcome)
def build_features(d: pd.DataFrame) -> pd.DataFrame:
    f = pd.DataFrame(index=d.index)
    f["limit"] = d["LIMIT_BAL"]
    f["age"] = d["AGE"]
    pay = d[PAY].clip(lower=-1)                      # -2 (no consumption) and -1 (paid in full) both mean "not late"
    f["late_now"] = pay["PAY_1"].clip(lower=0)       # months late on the most recent statement
    f["max_late"] = pay.clip(lower=0).max(axis=1)
    f["months_late"] = (pay > 0).sum(axis=1)
    w = np.array([6, 5, 4, 3, 2, 1], dtype=float)    # recency weights: last month counts most
    f["late_recency"] = (pay.clip(lower=0).values * w).sum(axis=1) / w.sum()
    f["late_trend"] = pay["PAY_1"] - pay["PAY_3"]    # getting worse (+) or better (−)
    bill = d[BILL]; pamt = d[PAMT]
    f["util_now"] = (bill["BILL_AMT1"] / d["LIMIT_BAL"]).clip(-1, 5)
    f["util_mean"] = (bill.mean(axis=1) / d["LIMIT_BAL"]).clip(-1, 5)
    f["bill_now"] = bill["BILL_AMT1"]
    f["bill_trend"] = (bill["BILL_AMT1"] - bill["BILL_AMT6"]) / (d["LIMIT_BAL"] + 1)
    prev_bill = bill[["BILL_AMT2", "BILL_AMT3", "BILL_AMT4", "BILL_AMT5", "BILL_AMT6"]].values
    paid = pamt[["PAY_AMT1", "PAY_AMT2", "PAY_AMT3", "PAY_AMT4", "PAY_AMT5"]].values
    ratio = np.where(prev_bill > 0, paid / np.maximum(prev_bill, 1), 1.0)   # share of last bill actually paid
    f["pay_ratio_mean"] = np.clip(ratio, 0, 2).mean(axis=1)
    f["pay_ratio_now"] = np.clip(ratio[:, 0], 0, 2)
    f["zero_pay_months"] = (pamt.values[:, :5] == 0).sum(axis=1)
    f["paid_total_6m"] = pamt.sum(axis=1)
    return f

X = build_features(df)
FEATURES = list(X.columns)
protected = df[["SEX", "AGE"]]

# ---------------------------------------------------------------- 3. split
X_tr, X_te, y_tr, y_te, p_tr, p_te = train_test_split(X, y, protected, test_size=0.2, stratify=y, random_state=SEED)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=SEED)

# ---------------------------------------------------------------- 4. models
lr = make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000, C=0.5))
gbm = GradientBoostingClassifier(n_estimators=220, learning_rate=0.05, max_depth=3, subsample=0.8, min_samples_leaf=40, random_state=SEED)

def cv_auc(model):
    p = cross_val_predict(model, X_tr, y_tr, cv=cv, method="predict_proba")[:, 1]
    return roc_auc_score(y_tr, p), average_precision_score(y_tr, p)

lr_cv = cv_auc(lr); gbm_cv = cv_auc(gbm)
lr.fit(X_tr, y_tr); gbm.fit(X_tr, y_tr)
p_lr = lr.predict_proba(X_te)[:, 1]; p_gb = gbm.predict_proba(X_te)[:, 1]
res = {
    "n_total": int(len(df)), "n_train": int(len(X_tr)), "n_test": int(len(X_te)), "base_rate": float(y.mean()),
    "logistic": {"cv_auc": lr_cv[0], "cv_pr_auc": lr_cv[1], "test_auc": roc_auc_score(y_te, p_lr), "test_pr_auc": average_precision_score(y_te, p_lr), "brier": brier_score_loss(y_te, p_lr)},
    "gbm": {"cv_auc": gbm_cv[0], "cv_pr_auc": gbm_cv[1], "test_auc": roc_auc_score(y_te, p_gb), "test_pr_auc": average_precision_score(y_te, p_gb), "brier": brier_score_loss(y_te, p_gb)},
}
print(json.dumps(res, indent=1))

# calibration (on the held-out set)
frac, mean_pred = calibration_curve(y_te, p_gb, n_bins=10, strategy="quantile")
fpr, tpr, _ = roc_curve(y_te, p_gb)
prec, rec, thr = precision_recall_curve(y_te, p_gb)

# ---------------------------------------------------------------- threshold economics
# chase policy: contact the top-k% by (probability × exposure). exposure = current bill; a contacted
# account that would have missed pays with probability RECOVER; each contact costs COST.
exposure = X_te["bill_now"].clip(lower=0).values
RECOVER, COST = 0.35, 250.0                                # NT$ — illustrative policy parameters, stated on the page
order = np.argsort(-(p_gb * exposure))
ks, recovered, cost, net = [], [], [], []
would_miss = y_te[order] == 1
exp_sorted = exposure[order]
cum_rec = np.cumsum(would_miss * exp_sorted * RECOVER)
for k in range(5, 101, 5):
    n = int(len(order) * k / 100)
    ks.append(k); recovered.append(float(cum_rec[n - 1])); cost.append(float(n * COST)); net.append(float(cum_rec[n - 1] - n * COST))
best_k = ks[int(np.argmax(net))]

# lift by decile of model score
dec = pd.qcut(pd.Series(p_gb).rank(method="first"), 10, labels=False)
lift = [float(y_te[dec == d].mean() / y.mean()) for d in range(9, -1, -1)]
miss_rate_by_decile = [float(y_te[dec == d].mean()) for d in range(9, -1, -1)]

# ---------------------------------------------------------------- 5. fairness: AUC and miss rate by group (SEX not a feature)
fair = {}
for label, mask in [("women", p_te["SEX"].values == 2), ("men", p_te["SEX"].values == 1), ("under_30", p_te["AGE"].values < 30), ("30_to_45", (p_te["AGE"].values >= 30) & (p_te["AGE"].values < 45)), ("45_plus", p_te["AGE"].values >= 45)]:
    fair[label] = {"n": int(mask.sum()), "auc": float(roc_auc_score(y_te[mask], p_gb[mask])), "actual_miss_rate": float(y_te[mask].mean()), "mean_score": float(p_gb[mask].mean())}

# feature importance (permutation-free: impurity importance from GBM) + LR coefficients
imp = sorted(zip(FEATURES, gbm.feature_importances_.tolist()), key=lambda t: -t[1])
coef = dict(zip(FEATURES, lr.named_steps["logisticregression"].coef_[0].tolist()))

# ---------------------------------------------------------------- 6. export
def tree_to_json(t):
    return {"left": t.children_left.tolist(), "right": t.children_right.tolist(), "feat": t.feature.tolist(),
            "thr": [round(float(v), 6) for v in t.threshold.tolist()], "val": [round(float(v[0][0]), 6) for v in t.value.tolist()]}
model_json = {
    "features": FEATURES, "learning_rate": gbm.learning_rate,
    "init": float(gbm.init_.class_prior_[1]) if hasattr(gbm.init_, "class_prior_") else None,
    "init_raw": float(np.log(y_tr.mean() / (1 - y_tr.mean()))),
    "trees": [tree_to_json(est[0].tree_) for est in gbm.estimators_],
    "feature_stats": {f: {"mean": float(X_tr[f].mean()), "std": float(X_tr[f].std()), "min": float(X_tr[f].min()), "max": float(X_tr[f].max())} for f in FEATURES},
    "note": "GradientBoostingClassifier, sklearn; raw score = init_raw + lr * sum(tree leaf values); p = sigmoid(raw)",
}
# sanity: reproduce sklearn's probabilities with the exported form
def predict_from_json(m, row):
    raw = m["init_raw"]
    for t in m["trees"]:
        node = 0
        while t["left"][node] != -1:
            node = t["left"][node] if row[t["feat"][node]] <= t["thr"][node] else t["right"][node]
        raw += m["learning_rate"] * t["val"][node]
    return 1 / (1 + math.exp(-raw))
check = np.array([predict_from_json(model_json, r) for r in X_te.values[:200]])
max_dev = float(np.abs(check - p_gb[:200]).max())
print("export reproduces sklearn within", max_dev)
assert max_dev < 1e-6, "exported model does not reproduce sklearn"

# reference sample for the in-browser rank (500 held-out scores with exposure), anonymised numbers only
ref_idx = rng.choice(len(X_te), 1000, replace=False)
reference = [{"p": round(float(p_gb[i]), 4), "e": float(exposure[i]), "y": int(y_te[i])} for i in ref_idx]

evaluation = {
    "results": res, "roc": {"fpr": [round(float(v), 4) for v in fpr[::25]], "tpr": [round(float(v), 4) for v in tpr[::25]]},
    "pr": {"recall": [round(float(v), 4) for v in rec[::25]], "precision": [round(float(v), 4) for v in prec[::25]]},
    "calibration": {"mean_pred": [round(float(v), 4) for v in mean_pred], "frac_pos": [round(float(v), 4) for v in frac]},
    "economics": {"k": ks, "recovered": recovered, "cost": cost, "net": net, "best_k": best_k, "recover_rate": RECOVER, "cost_per_contact": COST},
    "lift": lift, "miss_rate_by_decile": miss_rate_by_decile, "fairness": fair,
    "importance": imp, "lr_coef": coef, "reference": reference,
}
json.dump(model_json, open("model.json", "w"))
json.dump(evaluation, open("evaluation.json", "w"))
print("wrote model.json (%d trees) and evaluation.json" % len(model_json["trees"]))
print("best chase depth:", best_k, "% of accounts; net at best:", round(net[ks.index(best_k)]))
print("fairness:", json.dumps(fair, indent=1))
print("top features:", imp[:8])