"""Reproduce Rosa's published Test 1 downstream figure (Models A / B / C).

This is the exact downstream-model evaluation behind the published Test 1
numbers on the Rosa Customer Portal. 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 model trained on Rosa-debiased data
generalises to a fair ground truth better than one trained on the biased data,
without a fairness cost.

WHAT YOU NEED (all downloadable from the portal Data Validation page, plus two
files you produce by running Rosa on Test 1):

  Downloaded from the portal:
    test1-population.csv   5,000 rows, the FAIR ground truth (has fair_outcome
                           and gender). Rows 0-1,999 are the training region,
                           rows 2,000-4,999 are the held-out evaluation region.
    test1.csv              2,000 rows, the BIASED training input (has outcome).
    test1-inference.csv    3,000 rows, the BIASED held-out input (has outcome).

  Produced by running Rosa on Test 1 (the portal one-click flow, or the API):
    <train>_fair.csv       the debiased TRAINING output (Training job on
                           test1.csv, bias_columns=[gender],
                           ignore_columns=[outcome]).
    <inference>_fair.csv   the debiased HELD-OUT output (Inference job on
                           test1-inference.csv using that training_job_id).

THE THREE MODELS (all scored on the fair ground truth, held-out region only):

  Model A  a LogisticRegression trained on the BIASED data (test1.csv),
           scored on the biased held-out input. The status quo.
  Model B  the SAME model trained on Rosa's DEBIASED output, scored on Rosa's
           debiased held-out output. This is what Rosa delivers.
  Model C  a parity-thresholded version of Model A: post-processing that forces
           equal approval rates per group. The "fairness by quota" baseline.

Each model:
  * LogisticRegression(random_state=42, max_iter=1000)
  * StandardScaler on the numeric features, OneHotEncoder on the categoricals
    (StandardScaler matters: without it the solver does not converge and the
    numbers become unstable).
  * scored on the fair-outcome ground truth over the held-out region.
  * bias metric = |mean(pred | men) - mean(pred | women)| / std(predictions)
    (a standardised mean-difference in the model's predictions).

EXPECTED NUMBERS (R2 on the fair ground truth):

  Model A ~= 0.598   (biased data)           Rosa-free, so this is EXACT.
  Model C ~= 0.546   (parity baseline)       Rosa-free, so this is EXACT.
  Model B ~= 0.686   (Rosa-debiased)         the published 5-seed MEDIAN,
                                             range [0.666-0.723].

Model B is stochastic: FAN training uses random weight initialisation, so a
single Rosa run lands somewhere in that range, not on 0.686 exactly. The
published figure is the median over five independent Rosa runs. Models A and C
never touch Rosa, so they are deterministic and reproduce exactly. The ordering
C < A < B (parity worst, biased middle, Rosa best) is the "no trade-off" result:
Rosa matches parity's fairness while keeping more predictive quality.

USAGE:

  python reproduce-test1-downstream.py \
      --population test1-population.csv \
      --train test1.csv \
      --inference test1-inference.csv \
      --debiased-train  <train>_fair.csv \
      --debiased-inference <inference>_fair.csv

Omit --debiased-train / --debiased-inference to compute only Models A and C
(the exact, Rosa-free figures) plus instructions for producing the Rosa output.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import r2_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# --- Column inventory (must match the shipped Test 1 dataset) --------------

FINANCIAL_FEATURES = [
    "credit_score",
    "annual_income",
    "debt_to_income_ratio",
    "months_credit_history",
    "num_late_payments_12m",
    "num_late_payments_lifetime",
    "savings_balance",
    "credit_utilisation_pct",
    "employment_tenure_months",
    "monthly_net_income",
]
GENUINE_CATEGORICALS = [
    "property_ownership",
    "employment_status",
    "education_level",
]
GENDER_PROXIES_NUMERIC = [
    "parental_leave_months",
    "career_break_months",
    "part_time_hours_weekly",
    "commute_distance_km",
    "weekly_shopping_trips",
    "streaming_hours_weekly",
]
ETHNICITY_PROXY = "country_of_birth_uk"
AGE_PROXY = "pension_enrolled"

PROTECTED = "gender"
TARGET = "outcome"
FAIR_TARGET = "fair_outcome"

# The two binary proxies are 0/1 integers scaled with the financials; only the
# genuine multi-value categoricals are one-hot encoded. The protected gender
# column is never a model feature - bias can only reach the model via proxies.
NUMERIC_MODEL_COLS = FINANCIAL_FEATURES + GENDER_PROXIES_NUMERIC + [ETHNICITY_PROXY, AGE_PROXY]
CAT_MODEL_COLS = GENUINE_CATEGORICALS

_TRAIN_ROWS = 2000  # rows [_TRAIN_ROWS:] of the population are the eval region


def _make_scaled_pipeline(num_cols: list[str], cat_cols: list[str]) -> Pipeline:
    """LogisticRegression with StandardScaler on numerics, OHE on categoricals."""
    pre = ColumnTransformer(
        [
            ("num", StandardScaler(), num_cols),
            ("cat", OneHotEncoder(handle_unknown="ignore"), cat_cols),
        ]
    )
    return Pipeline([("pre", pre), ("clf", LogisticRegression(random_state=42, max_iter=1000))])


def _present_model_cols(df: pd.DataFrame) -> tuple[list[str], list[str]]:
    """(numeric, categorical) model columns present in `df`.

    Works identically on the biased frame (has gender) and Rosa's debiased
    output (gender dropped) because gender is not in the model universe.
    """
    num = [c for c in NUMERIC_MODEL_COLS if c in df.columns]
    cat = [c for c in CAT_MODEL_COLS if c in df.columns]
    return num, cat


def _fit_scaled_model(train_df: pd.DataFrame, label_col: str) -> tuple[Pipeline, list[str]]:
    num, cat = _present_model_cols(train_df)
    feature_cols = num + cat
    pipe = _make_scaled_pipeline(num, cat)
    pipe.fit(train_df[feature_cols], train_df[label_col].to_numpy())
    return pipe, feature_cols


def _parity_thresholds(
    pipe: Pipeline, train_df: pd.DataFrame, label_col: str, feature_cols: list[str]
) -> tuple[float, float]:
    """Grid-search the parity-achieving (female, male) threshold pair.

    For each female threshold, derive the male threshold (empirical quantile)
    that equalises training approval rates, then keep the parity-achieving pair
    with the highest training accuracy. "Fairness through post-processing" - the
    accuracy cost Rosa avoids.
    """
    proba = pipe.predict_proba(train_df[feature_cols])[:, 1]
    y = train_df[label_col].to_numpy()
    is_female = (train_df[PROTECTED] == "F").to_numpy()
    proba_f = proba[is_female]
    proba_m = proba[~is_female]

    best = {"acc": -1.0, "t_f": 0.5, "t_m": 0.5}
    for t_f in np.arange(0.05, 0.96, 0.01):
        target_rate = float((proba_f >= t_f).mean())
        if target_rate <= 0.0 or target_rate >= 1.0:
            continue
        t_m = float(np.quantile(proba_m, 1.0 - target_rate))
        preds = np.where(is_female, proba >= t_f, proba >= t_m).astype(int)
        acc = float((preds == y).mean())
        rate_f = float((proba[is_female] >= t_f).mean())
        rate_m = float((proba[~is_female] >= t_m).mean())
        if abs(rate_f - rate_m) <= 0.02 and acc > best["acc"]:
            best = {"acc": acc, "t_f": float(t_f), "t_m": t_m}
    return best["t_f"], best["t_m"]


def bias_metric(predictions: np.ndarray, groups: np.ndarray) -> float:
    """Effect size of group disparity: |mean(g0) - mean(g1)| / std(all).

    Groups encoded 0/1 (0=men, 1=women).
    """
    g0 = predictions[groups == 0].mean()
    g1 = predictions[groups == 1].mean()
    std = predictions.std()
    return float(abs(g0 - g1) / std) if std > 0 else 0.0


def main() -> None:
    ap = argparse.ArgumentParser(description="Reproduce Rosa's Test 1 A/B/C downstream figure.")
    ap.add_argument("--population", default="test1-population.csv", type=Path)
    ap.add_argument("--train", default="test1.csv", type=Path)
    ap.add_argument("--inference", default="test1-inference.csv", type=Path)
    ap.add_argument("--debiased-train", default=None, type=Path)
    ap.add_argument("--debiased-inference", default=None, type=Path)
    args = ap.parse_args()

    population = pd.read_csv(args.population)
    biased_train = pd.read_csv(args.train)
    biased_holdout = pd.read_csv(args.inference)

    # Fair ground truth over the held-out region (rows _TRAIN_ROWS:).
    fair_holdout = population.iloc[_TRAIN_ROWS:].reset_index(drop=True)
    fair_y = fair_holdout[FAIR_TARGET].to_numpy()
    fair_is_female = (fair_holdout[PROTECTED] == "F").to_numpy()
    groups = fair_is_female.astype(int)
    if len(fair_holdout) != len(biased_holdout):
        raise SystemExit(
            f"row mismatch: population held-out region has {len(fair_holdout)} rows but "
            f"{args.inference} has {len(biased_holdout)} - use the matching shipped files"
        )

    # --- Model A: trained on the biased data, scored on the fair ground truth.
    model_a, cols_a = _fit_scaled_model(biased_train, TARGET)
    proba_a = model_a.predict_proba(biased_holdout[cols_a])[:, 1]
    pred_a = model_a.predict(biased_holdout[cols_a])
    r2_a = float(r2_score(fair_y, proba_a))
    acc_a = float((pred_a == fair_y).mean())
    bias_a = bias_metric(pred_a, groups)

    # --- Model C: parity-thresholded version of Model A (post-processing).
    model_c, cols_c = _fit_scaled_model(biased_train, TARGET)
    t_f, t_m = _parity_thresholds(model_c, biased_train, TARGET, cols_c)
    proba_c_eval = model_c.predict_proba(biased_holdout[cols_c])[:, 1]
    pred_c = np.where(fair_is_female, proba_c_eval >= t_f, proba_c_eval >= t_m).astype(int)
    proba_c_adjusted = np.clip(
        np.where(fair_is_female, proba_c_eval + (0.5 - t_f), proba_c_eval + (0.5 - t_m)),
        0.0,
        1.0,
    )
    r2_c = float(r2_score(fair_y, proba_c_adjusted))
    acc_c = float((pred_c == fair_y).mean())
    bias_c = bias_metric(pred_c, groups)

    print(f"rows: train={len(biased_train)} held-out={len(biased_holdout)}")
    print("model   R2       accuracy   bias")
    print(f"A (biased)   {r2_a:+.3f}    {acc_a:.3f}     {bias_a:.3f}   [Rosa-free, exact ~0.598]")
    print(f"C (parity)   {r2_c:+.3f}    {acc_c:.3f}     {bias_c:.3f}   [Rosa-free, exact ~0.546]")

    if args.debiased_train is None or args.debiased_inference is None:
        print()
        print("Model B needs Rosa's debiased output. Run Rosa on Test 1:")
        print("  1. Training job on test1.csv (bias_columns=[gender], ignore_columns=[outcome])")
        print("     -> download <train>_fair.csv")
        print("  2. Inference job on test1-inference.csv with that training_job_id")
        print("     -> download <inference>_fair.csv")
        print("  then re-run with --debiased-train / --debiased-inference.")
        return

    # --- Model B: trained on Rosa's debiased output, scored on the fair truth.
    debiased_train = pd.read_csv(args.debiased_train)
    debiased_inference = pd.read_csv(args.debiased_inference)
    model_b, cols_b = _fit_scaled_model(debiased_train, TARGET)
    proba_b = model_b.predict_proba(debiased_inference[cols_b])[:, 1]
    pred_b = model_b.predict(debiased_inference[cols_b])
    r2_b = float(r2_score(fair_y, proba_b))
    acc_b = float((pred_b == fair_y).mean())
    bias_b = bias_metric(pred_b, groups)

    print(
        f"B (Rosa)     {r2_b:+.3f}    {acc_b:.3f}     {bias_b:.3f}   "
        f"[stochastic; published 5-seed median 0.686, range 0.666-0.723]"
    )
    print()
    ordering = " < ".join(m for _, m in sorted([(r2_c, "C"), (r2_a, "A"), (r2_b, "B")]))
    print(f"ordering by R2: {ordering}  (target C < A < B, the no-trade-off result)")


if __name__ == "__main__":
    main()