"""Claim 5: a COMPAS case study shows optimal fair classifiers under sufficiency
generally require GROUP-SPECIFIC decision thresholds rather than one universal
threshold.

Previously: "the COMPAS data experiment was not rerun."

Logic of the test. Under sufficiency the score is calibrated within each group,
so the utility-optimal *unconstrained* rule is a single threshold shared by all
groups. Imposing a fairness constraint on top changes that: the constrained
optimum generally needs different thresholds per group, and forcing one universal
threshold costs utility. Both are measured.
"""
import json, numpy as np, pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.calibration import CalibratedClassifierCV
from sklearn.model_selection import train_test_split

RES = {}


def load():
    df = pd.read_pickle("compas.pkl")
    df = df[df["race"].isin(["African-American", "Caucasian"])].copy()
    X = pd.get_dummies(df[["age", "priors_count", "juv_fel_count",
                           "juv_misd_count", "c_charge_degree", "sex"]],
                       drop_first=True).astype(float).values
    y = df["two_year_recid"].values.astype(int)
    a = (df["race"] == "African-American").values.astype(int)
    return X, y, a


def calibration_by_group(s, y, a, bins=8):
    out = {}
    edges = np.quantile(s, np.linspace(0, 1, bins+1))
    for g in (0, 1):
        m = a == g; rows = []
        for i in range(bins):
            sel = m & (s >= edges[i]) & (s <= edges[i+1])
            if sel.sum() < 25: continue
            rows.append({"bin": i, "mean_score": round(float(s[sel].mean()), 4),
                         "empirical_rate": round(float(y[sel].mean()), 4), "n": int(sel.sum())})
        gap = max(abs(r["mean_score"]-r["empirical_rate"]) for r in rows) if rows else None
        out["group%d" % g] = {"bins": rows, "max_abs_calibration_gap": round(gap, 4)}
    return out


def utility(s, y, a, tA, tC, cost_fp=1.0, benefit_tp=1.0):
    d = np.where(a == 1, s >= tA, s >= tC)
    tp = np.sum(d & (y == 1)); fp = np.sum(d & (y == 0))
    return float(benefit_tp*tp-cost_fp*fp)


def fpr(s, y, a, tA, tC, g):
    m = (a == g) & (y == 0)
    t = tA if g == 1 else tC
    return float(np.mean(s[m] >= t)) if m.sum() else 0.0


def main():
    X, y, a = load()
    Xtr, Xte, ytr, yte, atr, ate = train_test_split(X, y, a, test_size=0.5,
                                                    random_state=0, stratify=y)
    base = LogisticRegression(max_iter=2000)
    clf = CalibratedClassifierCV(base, method="isotonic", cv=5).fit(Xtr, ytr)
    s = clf.predict_proba(Xte)[:, 1]
    cal = calibration_by_group(s, yte, ate)
    print("  sufficiency check — max |score - empirical rate| per group: g0=%.4f g1=%.4f"
          % (cal["group0"]["max_abs_calibration_gap"], cal["group1"]["max_abs_calibration_gap"]), flush=True)

    grid = np.linspace(0.05, 0.95, 91)
    # unconstrained optimum: single threshold (sufficiency => shared)
    uu = [(utility(s, yte, ate, t, t), t) for t in grid]
    u_un, t_un = max(uu)
    print("  unconstrained optimum: single threshold t=%.3f  utility=%.1f" % (t_un, u_un), flush=True)

    rows = []
    for eps in (0.10, 0.05, 0.02, 0.01):
        # constrained: equal FPR across groups within eps
        best_pair, best_u = None, -1e18
        for tA in grid:
            fA = fpr(s, yte, ate, tA, tA, 1)
            for tC in grid:
                fC = fpr(s, yte, ate, tC, tC, 0)
                if abs(fA-fC) > eps: continue
                u = utility(s, yte, ate, tA, tC)
                if u > best_u: best_u, best_pair = u, (tA, tC)
        # constrained but forced to a single universal threshold
        bu_s, bt_s = -1e18, None
        for t in grid:
            if abs(fpr(s, yte, ate, t, t, 1)-fpr(s, yte, ate, t, t, 0)) > eps: continue
            u = utility(s, yte, ate, t, t)
            if u > bu_s: bu_s, bt_s = u, t
        rows.append({"eps_fpr_gap": eps,
                     "group_specific_thresholds": [round(best_pair[0], 3), round(best_pair[1], 3)],
                     "threshold_difference": round(abs(best_pair[0]-best_pair[1]), 3),
                     "utility_group_specific": best_u,
                     "best_single_threshold": (round(bt_s, 3) if bt_s is not None else None),
                     "utility_single": (bu_s if bt_s is not None else None),
                     "utility_loss_from_single": (round(best_u-bu_s, 2) if bt_s is not None else None)})
        print("  eps=%.2f  group-specific t=(%.3f, %.3f) diff=%.3f util=%.1f | best single t=%s util=%s loss=%s"
              % (eps, best_pair[0], best_pair[1], rows[-1]["threshold_difference"], best_u,
                 rows[-1]["best_single_threshold"], rows[-1]["utility_single"],
                 rows[-1]["utility_loss_from_single"]), flush=True)
    RES["claim5_compas"] = {"n_test": int(len(yte)), "calibration": cal,
        "unconstrained_single_threshold": round(t_un, 3), "unconstrained_utility": u_un,
        "rows": rows,
        "thresholds_differ_in_all": all(r["threshold_difference"] > 1e-9 for r in rows),
        "single_threshold_costs_utility_in_all": all(
            r["utility_loss_from_single"] is None or r["utility_loss_from_single"] > 0 for r in rows)}
    json.dump(RES, open("compas_results.json", "w"), indent=1)


if __name__ == "__main__":
    main(); print("DONE")