"""Reproduce Rosa's published full-raw COMPAS downstream figure.

This is the exact downstream-model evaluation behind the published COMPAS number
on the Rosa Customer Portal - a 68% median reduction [64-83%] in the racial
false-positive-rate gap. It is plain scikit-learn - it does NOT import or depend
on the Rosa engine. It is the evaluator's model, external to Rosa: you run it
yourself to check that a recidivism model built on Rosa-debiased COMPAS carries
much less racial disparity than one built on the raw data, for a small accuracy
cost.

WHAT YOU NEED (the two splits are on the portal Data Validation page; the two
debiased files you produce by running Rosa on those splits):

  Downloaded from the portal:
    compas-full-train.csv   ~4,645 rows, the RAW training split.
    compas-full-test.csv    ~4,648 rows, the RAW held-out test split.
    (stratified 50/50 by race + is_recid, seed 42, from the raw ProPublica data)

  Produced by running Rosa on COMPAS (the portal one-click flow, or the API):
    <train>_fair.csv        the debiased TRAINING output (Training job on
                            compas-full-train.csv, bias_columns=[race],
                            ignore_columns=[is_recid]).
    <test>_fair.csv         the debiased HELD-OUT output (Inference job on
                            compas-full-test.csv using that training_job_id).

THE TWO MODELS (predict is_recid; fit on the train split, scored leakage-free on
the held-out test split):

  Model A  a LogisticRegression on the RAW features (race excluded). The status
           quo - it still carries the racial disparity through proxy columns.
  Model B  the SAME model on Rosa's DEBIASED output (race removed from every
           column by Rosa). This is what Rosa delivers.

Each model:
  * StandardScaler + LogisticRegression, one-hot-encoded categoricals.
  * race is EXCLUDED from the features of both models. Rosa removes the racial
    signal from the OTHER columns; leaving race in would re-introduce it
    directly. The true race is re-attached from the raw test split by row index
    (Rosa preserves row order) only to MEASURE the disparity, never as a feature.

Two disparity measures, each the gap between African-American and Caucasian
defendants in the model's predictions:
  * FPR gap  = the false-positive-rate gap (among defendants who did NOT
               re-offend, the difference in the rate of being predicted to
               re-offend). This is ProPublica's own COMPAS metric.
  * PPR gap  = the predicted-positive-rate gap (demographic parity): the
               difference in the overall rate of being predicted to re-offend.

EXPECTED NUMBERS:

  FPR-gap reduction  ~= 68%  (published 5-seed MEDIAN, range 64-83%)
  PPR-gap reduction  ~= 69%  (published 5-seed MEDIAN, range 62-76%)
  accuracy cost      ~= 2-4 percentage points

COMPAS is high-variance run to run - FAN training is stochastic, so a SINGLE
Rosa output lands anywhere in (and a bit beyond) that range, not on 68% exactly.
The published figure is the MEDIAN over five independent Rosa runs. A single-run
reproduction well above the range (e.g. 90%+) is a lucky draw, not an error.

USAGE:

  python reproduce-compas-downstream.py \
      --raw-train  compas-full-train.csv \
      --raw-test   compas-full-test.csv \
      --debiased-train <train>_fair.csv \
      --debiased-test  <test>_fair.csv
"""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

TARGET = "is_recid"
PROT = "race"


def _features(df: pd.DataFrame, drop: list[str]) -> pd.DataFrame:
    """One-hot-encode the frame, dropping the given columns (target/protected)."""
    x = df.drop(columns=[c for c in drop if c in df.columns])
    return pd.get_dummies(x, drop_first=False).astype(float)


def _fit_eval(
    x_train: pd.DataFrame,
    y_train,
    x_test: pd.DataFrame,
    y_test,
    is_aa,
    is_cc,
) -> tuple[float, float, float]:
    """Fit StandardScaler + LogisticRegression; return (accuracy, PPR gap, FPR gap)."""
    x_train, x_test = x_train.align(x_test, join="left", axis=1, fill_value=0.0)
    model = make_pipeline(StandardScaler(), LogisticRegression(max_iter=3000))
    model.fit(x_train, y_train)
    pred = model.predict(x_test)
    acc = float((pred == y_test).mean())
    ppr_gap = float(abs(pred[is_aa].mean() - pred[is_cc].mean()))
    neg = y_test == 0  # false-positive rate is measured among true negatives
    fpr_gap = float(abs(pred[is_aa & neg].mean() - pred[is_cc & neg].mean()))
    return acc, ppr_gap, fpr_gap


def main() -> None:
    ap = argparse.ArgumentParser(description="Reproduce Rosa's full-raw COMPAS downstream figure.")
    ap.add_argument("--raw-train", default="compas-full-train.csv", type=Path)
    ap.add_argument("--raw-test", default="compas-full-test.csv", type=Path)
    ap.add_argument("--debiased-train", default=None, type=Path)
    ap.add_argument("--debiased-test", default=None, type=Path)
    args = ap.parse_args()

    raw_tr = pd.read_csv(args.raw_train)
    raw_te = pd.read_csv(args.raw_test)

    # True race, re-attached by row index for BOTH models (measurement only).
    race_te = raw_te[PROT].to_numpy()
    is_aa = race_te == "African-American"
    is_cc = race_te == "Caucasian"

    acc_a, ppr_a, fpr_a = _fit_eval(
        _features(raw_tr, [TARGET, PROT]),
        raw_tr[TARGET].to_numpy(),
        _features(raw_te, [TARGET, PROT]),
        raw_te[TARGET].to_numpy(),
        is_aa,
        is_cc,
    )

    print(
        f"rows: train={len(raw_tr)} test={len(raw_te)}  base is_recid rate={raw_te[TARGET].mean():.3f}"
    )
    print("model            accuracy   PPR gap   FPR gap")
    print(f"A (raw features)   {acc_a:.3f}     {ppr_a:.3f}     {fpr_a:.3f}")

    if args.debiased_train is None or args.debiased_test is None:
        print()
        print("Model B needs Rosa's debiased output. Run Rosa on COMPAS:")
        print("  1. Training job on compas-full-train.csv (bias_columns=[race],")
        print("     ignore_columns=[is_recid]) -> download <train>_fair.csv")
        print("  2. Inference job on compas-full-test.csv with that training_job_id")
        print("     -> download <test>_fair.csv")
        print("  then re-run with --debiased-train / --debiased-test.")
        return

    deb_tr = pd.read_csv(args.debiased_train)
    deb_te = pd.read_csv(args.debiased_test)
    if len(deb_te) != len(raw_te):
        raise SystemExit(
            f"row mismatch: debiased test has {len(deb_te)} rows but raw test has "
            f"{len(raw_te)} - race is re-attached by index, so they must align"
        )
    acc_b, ppr_b, fpr_b = _fit_eval(
        _features(deb_tr, [TARGET]),
        deb_tr[TARGET].to_numpy(),
        _features(deb_te, [TARGET]),
        deb_te[TARGET].to_numpy(),
        is_aa,
        is_cc,
    )
    print(f"B (Rosa-debiased)  {acc_b:.3f}     {ppr_b:.3f}     {fpr_b:.3f}")
    print()
    if fpr_a > 0 and ppr_a > 0:
        print(
            f"disparity reduction (single run):  "
            f"FPR gap {(1 - fpr_b / fpr_a) * 100:.0f}%   PPR gap {(1 - ppr_b / ppr_a) * 100:.0f}%"
        )
    print(f"accuracy cost (A - B): {acc_a - acc_b:+.3f}")
    print("published figure is the 5-seed MEDIAN: FPR gap 68% [64-83%], PPR gap 69% [62-76%]")


if __name__ == "__main__":
    main()