"""
The invoice that comes back — a billing-reversal model on real wholesale transactions.

Data: UCI "Online Retail II" (Chen, 2019) — 1,067,371 invoice lines from a UK online wholesaler, Dec 2009 – Dec 2011,
5,942 customers, 43 countries. Cancellations are their own invoices (numbers prefixed "C") with negative quantities.

Pipeline
  1. load + clean (pseudo stock codes, guest sales without a customer id, zero prices)
  2. reconciliation: tie every cancellation line back to the invoice line it reverses (same customer + product, most recent
     sale within 120 days) — the credited value and lag per invoice
  3. invoice table + features known when the invoice is raised; customer history strictly before that moment, and a prior
     reversal only counts if the cancellation had already happened (no lookahead)
  4. time split: train on invoices before 2011-05-01; test on May–Aug 2011 (reversals need time to arrive; data ends 9 Dec)
  5. models: logistic baseline and gradient boosting for P(reversed); baseline = the customer's own prior reversal rate
  6. evaluation: AUC / PR-AUC / Brier, calibration, lift, AUC by segment; check-before-dispatch economics
  7. export: trees → model.json; tables, curves, reference sample → evaluation.json
"""
import json, math
import numpy as np
import pandas as pd
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
from sklearn.calibration import calibration_curve

SEED = 42
rng = np.random.default_rng(SEED)
SPLIT = pd.Timestamp("2011-05-01"); TEST_END = pd.Timestamp("2011-09-01")
MATCH_WINDOW = pd.Timedelta(days=120)

# ---------------------------------------------------------------- 1. load + clean
d = pd.read_csv("data/online_retail_ii.csv", parse_dates=["InvoiceDate"])
d.columns = ["invoice", "stock", "desc", "qty", "date", "price", "cust", "country", "cancel", "value"]
d["stock"] = d["stock"].astype(str); d["invoice"] = d["invoice"].astype(str)
pseudo = ~d["stock"].str.match(r"^\d")
hyg = {"lines_raw": int(len(d)), "pseudo_lines": int(pseudo.sum()), "pseudo_codes": d.loc[pseudo, "stock"].value_counts().head(10).to_dict(),
       "guest_lines_share": float(d["cust"].isna().mean()), "zero_price_share": float((d["price"] <= 0).mean()),
       "invoices_all": int(d["invoice"].nunique()), "cancel_invoices_all": int(d.loc[d["cancel"], "invoice"].nunique()),
       "billed_value_all": float(d.loc[~d["cancel"], "value"].sum()), "credited_value_all": float(-d.loc[d["cancel"], "value"].sum())}
k = d[d["cust"].notna() & ~pseudo].copy(); k["cust"] = k["cust"].astype(int)
sales = k[~k["cancel"] & (k["qty"] > 0) & (k["price"] > 0)].copy()
canc = k[k["cancel"] & (k["qty"] < 0)].copy()

# ---------------------------------------------------------------- 2. reconciliation: cancellation line → the sale it reverses
s_sorted = sales.sort_values("date")[["cust", "stock", "date", "invoice", "qty", "price"]].rename(columns={"date": "sdate", "invoice": "sinv", "qty": "sqty", "price": "sprice"})
m = pd.merge_asof(canc.sort_values("date"), s_sorted, left_on="date", right_on="sdate", by=["cust", "stock"], direction="backward", tolerance=MATCH_WINDOW)
m["matched"] = m["sinv"].notna(); m["lag"] = (m["date"] - m["sdate"]).dt.days
recon = {"cancel_lines": int(len(m)), "matched_lines_share": float(m["matched"].mean()), "matched_value_share": float((-m.loc[m["matched"], "value"]).sum() / (-m["value"]).sum()),
         "lag_hist": [{"bin": lab, "n": int(((m.lag >= lo) & (m.lag < hi) & m.matched).sum())} for lab, lo, hi in [("same day", 0, 1), ("1–3 d", 1, 4), ("4–7 d", 4, 8), ("8–14 d", 8, 15), ("15–30 d", 15, 31), ("31–60 d", 31, 61), ("61–120 d", 61, 121)]],
         "lag_quantiles": {str(q): float(m.loc[m.matched, "lag"].quantile(q)) for q in [.25, .5, .75, .9]}}
rev = m[m["matched"]].groupby("sinv").agg(credit=("value", lambda v: float(-v.sum())), first_cancel=("date", "min"), cancel_lines=("stock", "size"))

# ---------------------------------------------------------------- 3. invoice table + features
first_buy = sales.groupby(["cust", "stock"])["date"].transform("min")
sales["is_new"] = (sales["date"] == first_buy).astype(int)
usual = sales[sales["date"] < SPLIT].groupby("stock")["price"].median()            # a product's usual price, from the training period only
sales["usual"] = sales["stock"].map(usual)
sales["ratio"] = np.where(sales["usual"].notna(), sales["price"] / sales["usual"].replace(0, np.nan), 1.0)
sales["ratio"] = sales["ratio"].fillna(1.0).clip(0.2, 5)
sales["line_value"] = sales["qty"] * sales["price"]
inv = sales.groupby("invoice").agg(cust=("cust", "first"), date=("date", "min"), country=("country", "first"), value=("line_value", "sum"), lines=("stock", "size"),
                                   products=("stock", "nunique"), qty=("qty", "sum"), max_line=("line_value", "max"), new_lines=("is_new", "sum"),
                                   ratio_mean=("ratio", "mean"), discount_lines=("ratio", lambda r: int((r < 0.9).sum())), premium_lines=("ratio", lambda r: int((r > 1.1).sum())),
                                   max_qty=("qty", "max"))
inv = inv.join(rev)
inv["reversed"] = inv["credit"].notna().astype(int)
inv["credit"] = inv["credit"].fillna(0.0)
inv = inv.sort_values("date").reset_index()

hist_cols = ["n_prior", "prior_reversed_rate", "prior_credit_share", "last_reversed", "days_since_last", "prior_value_mean", "prior_lines_mean", "tenure_days", "prior_reversed_count"]
H = np.full((len(inv), len(hist_cols)), -1.0)
for cust, g in inv.groupby("cust", sort=False):
    idx = g.index.values; dt = g["date"].values; fc = g["first_cancel"].values; val = g["value"].values; cr = g["credit"].values; ln = g["lines"].values
    for j, i in enumerate(idx):
        t = dt[j]; prior = dt < t; n = int(prior.sum())
        if n == 0: H[i, 0] = 0; continue
        known = prior & ~np.isnat(fc) & (fc < t)                       # a prior invoice counts as reversed only if the cancellation had already happened
        H[i] = [n, known.sum() / n, cr[known].sum() / max(val[prior].sum(), 1e-9), 1.0 if known[np.flatnonzero(prior)[-1]] else 0.0,
                (t - dt[prior].max()) / np.timedelta64(1, "D"), val[prior].mean(), ln[prior].mean(), (t - dt[prior].min()) / np.timedelta64(1, "D"), known.sum()]
hist = pd.DataFrame(H, columns=hist_cols, index=inv.index)

def build_features(v, h):
    f = pd.DataFrame(index=v.index)
    f["log_value"] = np.log1p(v["value"]); f["lines"] = v["lines"]; f["products"] = v["products"]; f["log_qty"] = np.log1p(v["qty"])
    f["max_line_share"] = v["max_line"] / v["value"]; f["new_share"] = v["new_lines"] / v["lines"]
    f["ratio_mean"] = v["ratio_mean"]; f["discount_share"] = v["discount_lines"] / v["lines"]; f["premium_share"] = v["premium_lines"] / v["lines"]
    f["max_qty"] = np.log1p(v["max_qty"])
    f["is_uk"] = (v["country"] == "United Kingdom").astype(int); f["is_eire"] = (v["country"] == "EIRE").astype(int)
    f["hour"] = v["date"].dt.hour; f["weekday"] = v["date"].dt.dayofweek; f["month"] = v["date"].dt.month
    for c in hist_cols: f[c] = h[c]
    return f
X_all = build_features(inv, hist); FEATURES = list(X_all.columns); y_all = inv["reversed"].values
tr = inv["date"] < SPLIT; te = (inv["date"] >= SPLIT) & (inv["date"] < TEST_END)
X_tr, X_te, y_tr, y_te = X_all[tr], X_all[te], y_all[tr], y_all[te]
print("train", int(tr.sum()), "test", int(te.sum()), "base", round(float(y_all[tr].mean()), 3), round(float(y_te.mean()), 3))

# ---------------------------------------------------------------- 4/5. models
lr = make_pipeline(StandardScaler(), LogisticRegression(max_iter=3000, C=0.5))
gbm = GradientBoostingClassifier(n_estimators=250, learning_rate=0.05, max_depth=3, subsample=0.8, min_samples_leaf=40, random_state=SEED)
lr.fit(X_tr, y_tr); gbm.fit(X_tr, y_tr)
p_lr, p_gb = lr.predict_proba(X_te)[:, 1], gbm.predict_proba(X_te)[:, 1]
GLOBAL = float(y_tr.mean())
base = np.where(X_te["n_prior"] > 0, X_te["prior_reversed_rate"], GLOBAL)
def metr(p): return {"auc": float(roc_auc_score(y_te, p)), "pr_auc": float(average_precision_score(y_te, p)), "brier": float(brier_score_loss(y_te, np.clip(p, 0, 1)))}
res = {"n_train": int(tr.sum()), "n_test": int(te.sum()), "split": str(SPLIT.date()), "test_end": str(TEST_END.date()), "base_rate_train": GLOBAL, "base_rate_test": float(y_te.mean()),
       "baseline_prior_rate": metr(base), "logistic": metr(p_lr), "gbm": metr(p_gb), "credited_share_of_billed": float(inv["credit"].sum() / inv["value"].sum()),
       "credit_share_when_reversed_median": float((inv.loc[inv.reversed == 1, "credit"] / inv.loc[inv.reversed == 1, "value"]).median()),
       "full_reversal_share": float(((inv.loc[inv.reversed == 1, "credit"] / inv.loc[inv.reversed == 1, "value"]) >= 0.95).mean())}
print(json.dumps(res, indent=1))

fpr, tpr, _ = roc_curve(y_te, p_gb)
frac, mean_pred = calibration_curve(y_te, p_gb, n_bins=10, strategy="quantile")
dec = pd.qcut(pd.Series(p_gb).rank(method="first"), 10, labels=False)
rate_by_decile = [float(y_te[dec == q].mean()) for q in range(9, -1, -1)]
credit_by_decile = [float(inv.loc[te, "credit"].values[dec == q].sum()) for q in range(9, -1, -1)]

# AUC by segment (test)
segs = {}
for label, mask in [("uk", X_te["is_uk"].values == 1), ("non_uk", X_te["is_uk"].values == 0), ("first_invoice", X_te["n_prior"].values == 0), ("1_to_5_prior", (X_te["n_prior"].values >= 1) & (X_te["n_prior"].values <= 5)),
                    ("6_to_20_prior", (X_te["n_prior"].values >= 6) & (X_te["n_prior"].values <= 20)), ("21_plus_prior", X_te["n_prior"].values > 20),
                    ("value_under_100", inv.loc[te, "value"].values < 100), ("value_100_to_1000", (inv.loc[te, "value"].values >= 100) & (inv.loc[te, "value"].values < 1000)), ("value_1000_plus", inv.loc[te, "value"].values >= 1000)]:
    if mask.sum() > 50 and 0 < y_te[mask].mean() < 1:
        segs[label] = {"n": int(mask.sum()), "auc": float(roc_auc_score(y_te[mask], p_gb[mask])), "actual_rate": float(y_te[mask].mean()), "mean_score": float(p_gb[mask].mean())}

# ---------------------------------------------------------------- 6. economics: check before dispatch
# rank test invoices by P(reversed) × value; a check costs COST; on an invoice that would have come back it recovers SAVE of the credited value
value_te = inv.loc[te, "value"].values; credit_te = inv.loc[te, "credit"].values
COST, SAVE = 4.0, 0.4
order = np.argsort(-(p_gb * value_te)); cum = np.cumsum(credit_te[order] * SAVE)
ks, saved, cost, net = [], [], [], []
for kk in range(5, 101, 5):
    n = int(len(order) * kk / 100); ks.append(kk); saved.append(float(cum[n - 1])); cost.append(float(n * COST)); net.append(float(cum[n - 1] - n * COST))
best_k = ks[int(np.argmax(net))]
ref_idx = rng.choice(len(y_te), 1000, replace=False)
reference = [{"p": round(float(p_gb[i]), 4), "v": round(float(value_te[i]), 2), "c": round(float(credit_te[i]), 2), "y": int(y_te[i])} for i in ref_idx]

# ---------------------------------------------------------------- descriptive tables (all invoices with a customer)
tables = {
    "by_value_band": [{"band": lab, "n": int(mk.sum()), "reversed_rate": float(inv.loc[mk, "reversed"].mean()), "credited": float(inv.loc[mk, "credit"].sum()), "billed": float(inv.loc[mk, "value"].sum())}
                      for lab, mk in [("< £100", inv.value < 100), ("£100–300", inv.value.between(100, 300)), ("£300–1,000", inv.value.between(300, 1000)), ("£1,000–3,000", inv.value.between(1000, 3000)), ("£3,000+", inv.value > 3000)]],
    "by_history": [{"band": lab, "n": int(mk.sum()), "reversed_rate": float(inv.loc[mk, "reversed"].mean())}
                   for lab, mk in [("first invoice", hist.n_prior == 0), ("never reversed before", (hist.n_prior > 0) & (hist.prior_reversed_rate == 0)), ("0–10% reversed", (hist.prior_reversed_rate > 0) & (hist.prior_reversed_rate <= .1)),
                                   ("10–30% reversed", (hist.prior_reversed_rate > .1) & (hist.prior_reversed_rate <= .3)), ("30%+ reversed", hist.prior_reversed_rate > .3)]],
    "by_month": [{"month": str(kq), "n": int(len(g)), "reversed_rate": float(g.reversed.mean())} for kq, g in inv.groupby(inv.date.dt.to_period("M"))],
    "by_country": [{"country": c, "n": int(len(g)), "reversed_rate": float(g.reversed.mean()), "credited": float(g.credit.sum())} for c, g in inv.groupby("country") if len(g) >= 150],
    "by_new_share": [{"band": lab, "n": int(mk.sum()), "reversed_rate": float(inv.loc[mk, "reversed"].mean())}
                     for lab, mk in [("0–25% new", X_all.new_share <= .25), ("25–50%", (X_all.new_share > .25) & (X_all.new_share <= .5)), ("50–75%", (X_all.new_share > .5) & (X_all.new_share <= .75)), ("75–100% new", X_all.new_share > .75)]],
    "credit_share_hist": [{"bin": lab, "n": int(mk.sum())} for lab, mk in [(l, ((inv.credit / inv.value) >= lo) & ((inv.credit / inv.value) < hi) & (inv.reversed == 1)) for l, lo, hi in
                                                                       [("< 2%", 0, .02), ("2–5%", .02, .05), ("5–10%", .05, .1), ("10–25%", .1, .25), ("25–50%", .25, .5), ("50–95%", .5, .95), ("full", .95, 10)]]],
    "top_products": [{"stock": s, "desc": str(g["desc"].iloc[0])[:40], "credited": float(-g["value"].sum()), "lines": int(len(g))} for s, g in m[m.matched].groupby("stock") if True][:0],
}
tp = m[m.matched].groupby("stock").agg(credited=("value", lambda v: float(-v.sum())), lines=("stock", "size"), desc=("desc", "first")).sort_values("credited", ascending=False).head(10)
tables["top_products"] = [{"stock": s, "desc": str(r.desc)[:40], "credited": float(r.credited), "lines": int(r.lines)} for s, r in tp.iterrows()]
imp = sorted(zip(FEATURES, gbm.feature_importances_.tolist()), key=lambda t: -t[1])
coef = dict(zip(FEATURES, lr.named_steps["logisticregression"].coef_[0].tolist()))

# ---------------------------------------------------------------- 7. export
def tree_to_json(t):
    return {"left": t.children_left.tolist(), "right": t.children_right.tolist(), "feat": t.feature.tolist(), "thr": [float(v) 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_raw": float(np.log(GLOBAL / (1 - GLOBAL))), "trees": [tree_to_json(e[0].tree_) for e in gbm.estimators_],
              "global_rate": GLOBAL, "credit_share_when_reversed": float((inv.loc[(inv.reversed == 1) & tr, "credit"] / inv.loc[(inv.reversed == 1) & tr, "value"]).mean()),
              "note": "GradientBoostingClassifier; P(reversed) = sigmoid(init_raw + lr * sum(leaf)). History features are -1 (n_prior = 0) for a customer's first invoice."}
def predict_from_json(mj, row):
    raw = mj["init_raw"]
    for t in mj["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 += mj["learning_rate"] * t["val"][node]
    return 1 / (1 + math.exp(-raw))
dev = float(np.abs(np.array([predict_from_json(model_json, r) for r in X_te.values[:300]]) - p_gb[:300]).max()); print("export dev", dev); assert dev < 1e-6
presets = {}
for name, cond in [("clean_regular", (X_te.n_prior >= 15) & (X_te.prior_reversed_rate == 0) & (inv.loc[te, "value"] > 250)), ("serial_returner", (X_te.n_prior >= 5) & (X_te.prior_reversed_rate >= .4) & (inv.loc[te, "value"] > 400)), ("big_first_order", (X_te.n_prior == 0) & (inv.loc[te, "value"] > 800))]:
    i = X_te[cond].index[0]; pos = list(X_te.index).index(i)
    presets[name] = {f: float(X_te.loc[i, f]) for f in FEATURES}; presets[name].update({"p": float(p_gb[pos]), "value": float(inv.loc[i, "value"]), "actual": int(y_te[pos])})

evaluation = {"results": res, "hygiene": hyg, "reconciliation": recon, "roc": {"fpr": [round(float(v), 4) for v in fpr[::10]], "tpr": [round(float(v), 4) for v in tpr[::10]]},
              "calibration": {"mean_pred": [round(float(v), 4) for v in mean_pred], "frac_pos": [round(float(v), 4) for v in frac]}, "rate_by_decile": rate_by_decile, "credit_by_decile": credit_by_decile,
              "segments": segs, "economics": {"k": ks, "saved": saved, "cost": cost, "net": net, "best_k": best_k, "cost_per_check": COST, "save_rate": SAVE}, "tables": tables,
              "importance": imp, "lr_coef": coef, "reference": reference, "presets": presets,
              "totals": {"invoices": int(len(inv)), "customers": int(inv.cust.nunique()), "reversed_invoices": int(inv.reversed.sum()), "billed": float(inv.value.sum()), "credited": float(inv.credit.sum()), "customers_with_reversal": float(inv.groupby("cust").reversed.max().mean())}}
json.dump(model_json, open("model.json", "w")); json.dump(evaluation, open("evaluation.json", "w"))
print("wrote model.json (%d trees), evaluation.json" % len(model_json["trees"]))
print("segments", json.dumps(segs, indent=1)); print("importance", imp[:8]); print("econ best_k", best_k, [round(n) for n in net][:8]); print("decile", [round(v, 3) for v in rate_by_decile])
print("by_history", tables["by_history"]); print("by_new_share", tables["by_new_share"]); print("top products", tables["top_products"][:5]); print("totals", evaluation["totals"])