"""Independently confirm Rosa's claim on a synthetic test dataset (Models A/B/C).

This is the published third-party validator for datasets built by the companion
generator (`synthetic_test_generator.py`). It is plain scikit-learn - it imports
NOTHING from the Rosa engine or wrapper. A customer, regulator or auditor runs it
to check, without trusting Rosa's internals, that a model trained on Rosa-debiased
data generalises to the FAIR ground truth better than a model trained on the
biased data, with no meaningful accuracy loss. This script is ONLY for a dataset
built by the companion synthetic_test_generator.py (it expects a `fair_outcome`
column and a `protected` bias column); to reproduce the portal's own Test 1 or
COMPAS figures use reproduce-test1-downstream.py or reproduce-compas-downstream.py
instead - they are dataset-specific and not interchangeable.

The synthetic dataset carries two outcome columns:
  outcome        the BIASED label. Model A trains on it. Passed to Rosa as an
                 ignore_column, so it survives into the debiased output.
  fair_outcome   the FAIR ground truth. The scoring key ONLY. Submitted to Rosa
                 only as an ignored pass-through (never trained on, never
                 debiased) and excluded from every model, so it survives into the
                 debiased output where this script reads it back.
The protected attribute is the bias_column: Rosa DROPS it from the debiased
output, so the true group labels are re-attached from the raw holdout by row
index (Rosa preserves row order) purely to MEASURE disparity - never as a feature.

Rosa preserves row ORDER but does not guarantee output COLUMN order, so every
model column is selected by NAME (never by position), and the re-attach by index
is guarded: a unique ignored `audit_row_id` is checked row-for-row between the raw
and debiased frames (training and holdout) before any metric, so a reorder, shuffle,
truncation or wrong-pair fails loudly instead of silently mis-scoring.

WHAT THE ROW-ID GATE PROVES, AND WHAT IT DOES NOT
  It proves ALIGNMENT: that the file you passed has the same rows, in the same
  order, as the raw file it claims to correspond to.
  It does NOT prove LINEAGE. Two independent Rosa runs over the same raw rows
  produce the SAME id sequence, so pairing run 1's training output with run 2's
  inference output passes this gate while being scientifically meaningless. That
  is not hypothetical: an external auditor did exactly that, Model B's R2 fell
  from 0.714 to 0.500, and an earlier version of this script still exited zero.
  To prove lineage, pass --training-manifest and --inference-manifest: the
  inference Run Manifest carries `training_job_id` (schema 1.3 onward), and this
  script checks it names the training job that produced the other file. Without
  those two arguments, lineage is reported as UNVERIFIED and you are trusting
  your own file handling.

THREE STRUCTURAL CHECKS ALWAYS FAIL CLOSED (they are integrity, not science):
  the row-id alignment gate; every ignored pass-through column matching the raw
  frame row-for-row (a flipped ignored outcome otherwise sails through and
  silently produces a nonsensical R2); and the protected attribute being ABSENT
  from a debiased output (Rosa drops it - a file that still carries it is not a
  Rosa fair output, and merely excluding it from the features would hide that).

THE UTILITY GATE IS THREE-STATE, AND ROW COUNT IS A PRECONDITION, NOT A RESULT
  Below 1,500 eligible training rows this script renders NO VERDICT on the C < A < B
  ordering rather than a pass or a fail, because at that size the ordering is close
  to a coin flip: measured over seven arms, 5 failures in 20 runs at ~1,200 training
  rows against 0 in 40 runs at 1,500 or more. At or above the floor it reports the
  actual held-out result. It never infers a pass from the row count in either
  direction - a count above the floor makes the result READABLE, it does not make it
  a pass. Under --enforce-gates a NO VERDICT exits non-zero for the same reason a
  failure does: a script used as an approval gate must decline to approve when it
  cannot render a verdict.

  One invocation is ONE Rosa run. The published protocol is ten independent runs
  aggregated as a mean; Rosa exposes no training seed, so the runs are the
  measurement.

FOUR PER-RUN FACTS ARE PRINTED, RATHER THAN A CELL THRESHOLD
  Eligible training rows; protected-group counts; protected x outcome joint counts
  (with the smallest cell); and whether Rosa accepted or declined. Facts rather than
  a threshold because minimum cell count alone predicts neither outcome - a balanced
  set with a 221-row minimum cell was accepted and failed the utility gate 30% of the
  time, while an imbalanced set with a 225-row minimum cell was declined outright.
  Three of the four are also on the Run Manifest; the JOINT counts are not, and
  cannot be, because Rosa is never told which column is the outcome.

  Bias removal and utility preservation are INDEPENDENT AXES. In the failing arms the
  manifest residual read "not detected" on every run, including the runs that failed
  the utility gate, so nothing in the report or the manifest tells you which side of
  the gate a run landed on. That is precisely why this held-out acceptance test is
  worth running, and why it should be pre-registered for a consequential deployment.

Three independent sources of randomness sit behind these numbers, kept distinct on purpose:
  * Rosa's FAN training - stochastic, with no seed exposed, which is why Model B
    varies run to run (run 10 independent Rosa runs and take the mean);
  * the downstream MODEL seed - LogisticRegression(random_state=42), deterministic;
  * the recovery-probe SPLIT seed - train_test_split(random_state=1), deterministic.

Model C's R2 is computed on a CONSTRUCTED score (Model A's probability shifted per
group by 0.5 - threshold) so the quota baseline has a comparable continuous output;
it is a deliberately naive baseline, NOT a calibrated probability, so no Brier score
is reported for C (Brier is shown only for A and B, whose predict_proba is real).

WHAT YOU NEED:
  Produced by the generator:
    <biased-train>        raw biased training CSV (protected + features + outcome
                          + fair_outcome).
    <raw-holdout>         raw biased held-out CSV (same schema) - the inference
                          input.
  Produced by running Rosa on those two files:
    <debiased-train>      Rosa TRAINING output on <biased-train>
                          (bias_columns=[protected], ignore_columns=[outcome,
                          fair_outcome, audit_row_id]) - the debiased features, no protected.
    <debiased-holdout>    Rosa INFERENCE output on <raw-holdout> using that
                          training_job_id - debiased features, no protected.

THE THREE MODELS (all scored on the held-out fair_outcome):
  Model A  LogisticRegression on the BIASED data, scored on the raw holdout.
  Model B  the SAME model on Rosa's DEBIASED output, scored on the debiased
           holdout. This is what Rosa delivers.
  Model C  Model A with per-group parity thresholds (equal approval rates) - the
           "fairness by quota" post-processing baseline.
Each: StandardScaler on numerics + OneHotEncoder on categoricals + LogisticRegression.
Report R2 (the decisive generalisation gate), group bias, accuracy; expect the
no-trade-off ordering C < A < B on R2. Plus the non-linear protected-attribute
recovery (an MLP vs a logistic model) on the biased vs debiased features: it must
DROP after debiasing, proving Rosa removed recoverability a correlation check misses.

USAGE:
  python reproduce-synthetic-downstream.py \
      --biased-train    train.csv \
      --raw-holdout     holdout.csv \
      --debiased-train  train_fair.csv \
      --debiased-holdout holdout_fair.csv \
      [--training-manifest training.json --inference-manifest inference.json] \
      [--enforce-gates]

Omit the two debiased files to compute only Models A and C (Rosa-free).

--enforce-gates makes the prespecified SCIENTIFIC gates (the C < A < B ordering
and the fall in protected-attribute recoverability) exit non-zero when they fail.
The default is to report them, because a failed ordering on your own data can be
a real scientific result rather than an error - but if you are using this script
as an approval gate, pass --enforce-gates or a failure is only a line of text.
"""

from __future__ import annotations

import argparse
import json
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.exceptions import ConvergenceWarning
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    brier_score_loss,
    r2_score,
    roc_auc_score,
)
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import LabelEncoder, OneHotEncoder, StandardScaler

TARGET = "outcome"
FAIR_TARGET = "fair_outcome"
ROW_ID = "audit_row_id"

# The measured floor for the published utility claim, and the run count the
# published protocol uses. Below the floor this script renders NO VERDICT rather
# than a pass or a fail: the ordering is genuinely close to a coin flip there, so
# a single small run is an observation, never a result.
UTILITY_CLAIM_MIN_TRAINING_ROWS = 1500
PUBLISHED_PROTOCOL_RUNS = 10

# The published floor and its bound were measured on BINARY benchmarks. On a
# multi-class protected attribute the ordering line below is still printed, and
# the [PASS]/[FAIL] token stays (the portal e2e suite and the kit README quote
# it), but it is a RESULT on this data, not a verdict against a published claim -
# so the token is followed by this scope line. The per-class reading of the floor
# was tested and refuted (Fisher p = 1.00), so no per-group floor is enforced.
MULTICLASS_SCOPE_NOTE = (
    "               scope: the published 1,500-row floor and its 7.5% bound were measured on\n"
    "               binary benchmarks; on multi-class datasets the pass rate varied from 0/10\n"
    "               to 10/10 across datasets at identical settings, so this line is a result\n"
    "               on this data, not a verdict against a published claim."
)

# Rosa trains on a sample of at most about this many rows (drawn roughly in
# proportion to the groups); a training input above it comes back shorter than
# it went in, and the row-identity gate then fails closed on the count. Quoted
# in the explanation printed for that failure.
TRAINING_SAMPLE_ROWS = 10_000

UTILITY_FLOOR_NOTE = (
    "In the named synthetic benchmark, on Rosa 1.48.0 and the stated held-out A/B/C\n"
    "protocol, we observed 5 utility-ordering failures in 20 runs at approximately\n"
    "1,200 training rows and no failures in 40 runs at 1,500 or more. We have not\n"
    'characterised the transition between those ranges. Accordingly, the "no\n'
    'fairness/accuracy trade-off" claim is made only for runs with at least 1,500\n'
    "eligible training rows, subject to the stated acceptance test."
)


def _explain_sampled_training(
    biased: pd.DataFrame, deb_train: pd.DataFrame, training_manifest: Path | None
) -> None:
    """Fail closed, WITH the reason, when the training output is shorter than its input.

    Rosa samples a training input above ~10,000 rows before debiasing, so the
    debiased training output has fewer rows than the file that was submitted. The
    row-identity gate is right to refuse to score that (it cannot align a sampled
    output to the full input), but without this the customer sees only "row counts
    must match" and has no way to know why. The output row count is compared
    against the MANIFEST's `row_count` - never against the input file, because a
    large training output is legitimately shorter than its input.
    """
    n_in, n_out = len(biased), len(deb_train)
    if n_out >= n_in:
        return  # not the sampled case; the row-identity gate handles everything else
    manifest_rows = sampled_from = None
    if training_manifest is not None:
        try:
            m = _read_manifest(training_manifest)
            manifest_rows, sampled_from = m.get("row_count"), m.get("sampled_from_row_count")
        except (OSError, ValueError):
            pass
    if sampled_from and manifest_rows == n_out:
        why = (
            f"the training manifest records that Rosa SAMPLED the input: row_count={n_out} "
            f"(the rows debiased) from sampled_from_row_count={sampled_from} submitted"
        )
    elif training_manifest is not None:
        why = (
            f"the training manifest does not confirm sampling (row_count={manifest_rows}, "
            f"sampled_from_row_count={sampled_from}), so this may be a truncated or wrong file"
        )
    else:
        why = (
            f"a training input above about {TRAINING_SAMPLE_ROWS:,} rows is sampled by Rosa "
            "before debiasing (pass --training-manifest to confirm: its row_count is the rows "
            "debiased and sampled_from_row_count the rows submitted)"
        )
    raise SystemExit(
        f"training output: {n_out} rows vs {n_in} submitted - refusing to score. Reason: {why}.\n"
        "A sampled training output cannot be aligned row-for-row to the full input, so this "
        "validator fails closed on it. To score a run, submit a training split of fewer than "
        f"{TRAINING_SAMPLE_ROWS:,} rows (no sampling), keep the held-out split for inference, "
        "and re-run; the full file is still debiased end to end at the inference step."
    )


def _verify_row_ids(ref: pd.DataFrame, other: pd.DataFrame, which: str) -> None:
    """Fail closed unless ``other`` carries the ignored ``audit_row_id`` key
    row-for-row identical to ``ref`` - catching a reorder, within-label shuffle,
    wrong-pair, duplicate or truncation before any metric or by-index re-attach.
    """
    for tag, df in (("expected", ref), ("actual", other)):
        if ROW_ID not in df.columns:
            raise SystemExit(f"{which}: the {tag} frame is missing the '{ROW_ID}' column")
    a = ref[ROW_ID].astype(str).to_numpy()
    b = other[ROW_ID].astype(str).to_numpy()
    if len(a) != len(set(a)):
        raise SystemExit(f"{which}: '{ROW_ID}' is not unique")
    if len(b) != len(a):
        raise SystemExit(f"{which}: {len(b)} rows vs {len(a)} - row counts must match")
    if not np.array_equal(a, b):
        raise SystemExit(
            f"{which}: '{ROW_ID}' does not match row-for-row - reordered, shuffled, truncated "
            "or wrong file; refusing to score"
        )


def _verify_passthrough(
    ref: pd.DataFrame, other: pd.DataFrame, cols: list[str], which: str
) -> None:
    """Fail closed unless every IGNORED pass-through column survived unchanged.

    Rosa carries ignore_columns through untouched, so any difference means the file
    is not what it claims to be. Without this, a flipped ignored outcome column is
    invisible: it is excluded from the model features, so it never trips the
    feature checks, and it silently poisons every metric that scores against it.
    """
    for col in cols:
        if col not in other.columns:
            raise SystemExit(
                f"{which}: ignored pass-through column '{col}' is missing from the "
                f"debiased output - Rosa returns ignore_columns unchanged, so this "
                f"file is not a Rosa output for this input"
            )
        a = ref[col].reset_index(drop=True)
        b = other[col].reset_index(drop=True)
        if not a.equals(b):
            n_diff = int((a != b).sum())
            raise SystemExit(
                f"{which}: ignored pass-through column '{col}' differs from the raw "
                f"frame in {n_diff} row(s). Rosa passes ignored columns through "
                f"unchanged, so this file has been altered or is the wrong file; "
                f"refusing to score"
            )


def _verify_protected_absent(df: pd.DataFrame, protected: str, which: str) -> None:
    """Fail closed if a purported fair output still carries the protected column.

    Rosa DROPS the bias column from its debiased output. Merely excluding it from
    the model features would let a file that still contains it score as if it were
    debiased, which is the shape of a fairwashed result.
    """
    if protected in df.columns:
        raise SystemExit(
            f"{which}: the protected column '{protected}' is PRESENT in a file "
            f"passed as a Rosa debiased output. Rosa drops it; refusing to score."
        )


def _read_manifest(path: Path) -> dict:
    """Read a Run Manifest in either shape Rosa hands one out.

    `GET /v1/jobs/{job_id}/manifest` wraps it in the API's standard envelope,
    `{"status": "ok", "data": {...}}`; a manifest saved as the bare object is
    accepted too. An API error response is refused by name rather than read as
    a manifest with no fields.
    """
    doc = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(doc, dict):
        raise ValueError(f"{path.name} is not a JSON object")
    if doc.get("status") == "error" and isinstance(doc.get("error"), dict):
        code = doc["error"].get("code")
        raise ValueError(f"{path.name} is an API error response ({code}), not a Run Manifest")
    if "job_id" not in doc and isinstance(doc.get("data"), dict):
        return doc["data"]
    return doc


def _verify_lineage(training_manifest: Path | None, inference_manifest: Path | None) -> str:
    """Bind the inference output to the training run that produced its model.

    Returns a one-line human-readable status. The row-id gate proves alignment
    only - two independent Rosa runs over the same rows produce identical id
    sequences - so lineage has to be DECLARED, and it is: an inference Run
    Manifest carries `training_job_id` (manifest schema 1.3 onward).
    """
    if training_manifest is None and inference_manifest is None:
        return (
            "lineage: UNVERIFIED (pass --training-manifest and --inference-manifest "
            "to prove the debiased holdout came from this training run; the row-id "
            "gate proves alignment only)"
        )
    if training_manifest is None or inference_manifest is None:
        raise SystemExit(
            "lineage: pass BOTH --training-manifest and --inference-manifest, or neither"
        )
    try:
        train = _read_manifest(training_manifest)
        infer = _read_manifest(inference_manifest)
    except (OSError, ValueError) as exc:
        raise SystemExit(f"lineage: could not read a Run Manifest: {exc}") from exc

    train_id = train.get("job_id")
    declared = infer.get("training_job_id")
    if not train_id:
        raise SystemExit("lineage: the training manifest carries no `job_id`")
    if declared is None:
        raise SystemExit(
            "lineage: the inference manifest carries no `training_job_id`. Manifests "
            "written before schema 1.3 have no parent field, so lineage cannot be "
            "proved for this run - re-run the inference job to get a manifest that "
            "declares its parent."
        )
    if declared != train_id:
        raise SystemExit(
            f"lineage: the inference manifest declares training_job_id "
            f"'{declared}' but the training manifest is job '{train_id}'. These two "
            f"outputs are from DIFFERENT Rosa runs; refusing to score."
        )
    return f"lineage: VERIFIED (inference declares training_job_id {declared})"


def _job_status(training_manifest: Path | None) -> str:
    """Whether Rosa ACCEPTED or DECLINED the training run, read from its manifest.

    `declined_no_bias` is a terminal decision with a full evidence manifest, not a
    failure - and a decline is a controlled non-result, NOT evidence that the data
    are fair. Reported rather than inferred, because the absence of a debiased file
    could equally mean the job was never run.
    """
    if training_manifest is None:
        return "not supplied (pass --training-manifest to record it)"
    try:
        status = _read_manifest(training_manifest).get("job_status")
    except (OSError, ValueError):
        return "unreadable training manifest"
    if status == "declined_no_bias":
        return (
            "DECLINED (declined_no_bias) - Rosa found nothing above chance and stopped. "
            "A controlled non-result, not evidence that the data are fair."
        )
    if status == "declined_insufficient_support":
        return (
            "DECLINED (declined_insufficient_support) - a group of the bias column had too "
            "few rows for the detection test to run, so no test was run. Not a finding of "
            "no bias."
        )
    return f"ACCEPTED ({status})" if status else "no job_status in the training manifest"


def _print_run_facts(biased: pd.DataFrame, protected: str, training_manifest: Path | None) -> int:
    """Print the four per-run facts and return the eligible training-row count.

    Published as FACTS rather than reduced to a threshold, because our own
    evidence shows minimum cell count alone predicts neither outcome: a balanced
    set with a 221-row minimum cell was accepted and failed the utility gate 30%
    of the time, while an imbalanced set with a 225-row minimum cell was declined
    outright. Three of these four are also on the Run Manifest; the protected x
    outcome JOINT counts are not, and cannot be, because Rosa is never told which
    column is the outcome - the label sits anonymously inside `ignore_columns`.
    """
    eligible = len(biased)
    print("\nper-run facts (report these alongside any figure from this run):")
    print(f"  eligible training rows                  {eligible}")

    group_counts = biased[protected].value_counts().sort_index()
    print(
        "  protected-group counts                  "
        + ", ".join(f"{k}={int(v)}" for k, v in group_counts.items())
    )

    if TARGET in biased.columns:
        joint = pd.crosstab(biased[protected], biased[TARGET])
        cells = ", ".join(
            f"{g}x{o}={int(joint.loc[g, o])}" for g in joint.index for o in joint.columns
        )
        print(f"  protected x outcome cells               {cells}")
        print(f"  smallest cell                           {int(joint.to_numpy().min())}")
    else:
        print(f"  protected x outcome cells               n/a ('{TARGET}' not in the frame)")

    print(f"  Rosa accepted or declined               {_job_status(training_manifest)}")
    return eligible


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

    Excludes the protected attribute, both outcome columns and the row-id key.
    Categoricals are the non-numeric feature columns; identical logic works on the
    biased frame (has protected) and the debiased output (protected dropped).
    """
    excluded = {protected, TARGET, FAIR_TARGET, ROW_ID}
    feats = [c for c in df.columns if c not in excluded]
    cat = [c for c in feats if df[c].dtype == object]
    num = [c for c in feats if c not in cat]
    return num, cat


def _pipeline(num: list[str], cat: list[str]) -> Pipeline:
    pre = ColumnTransformer(
        [
            ("num", StandardScaler(), num),
            ("cat", OneHotEncoder(handle_unknown="ignore"), cat),
        ]
    )
    return Pipeline([("pre", pre), ("clf", LogisticRegression(random_state=42, max_iter=2000))])


def _fit(train_df: pd.DataFrame, label: str, protected: str) -> tuple[Pipeline, list[str]]:
    num, cat = _feature_cols(train_df, protected)
    cols = num + cat
    pipe = _pipeline(num, cat)
    pipe.fit(train_df[cols], train_df[label].to_numpy())
    return pipe, cols


def group_bias(pred: np.ndarray, groups: np.ndarray) -> float:
    """Max-min gap in mean prediction across the protected groups (any K)."""
    means = pd.Series(pred).groupby(np.asarray(groups)).mean()
    return float(means.max() - means.min())


def _parity(
    pipe: Pipeline, train_df: pd.DataFrame, cols: list[str], groups_train: np.ndarray
) -> dict[str, float]:
    """Per-group thresholds equalising approval rate, best training accuracy.

    Grid a global target approval rate; per group, threshold = the empirical
    quantile that hits it; keep the target with the highest training accuracy.
    Generalises the binary parity baseline to any number of protected groups.
    """
    proba = pipe.predict_proba(train_df[cols])[:, 1]
    y = train_df[TARGET].to_numpy()
    labels = np.unique(groups_train)
    best = {"acc": -1.0, "target": 0.5}
    thresholds: dict[str, float] = {}
    for target in np.arange(0.10, 0.91, 0.02):
        thr = {g: float(np.quantile(proba[groups_train == g], 1.0 - target)) for g in labels}
        pred = np.array([proba[i] >= thr[groups_train[i]] for i in range(len(proba))], dtype=int)
        acc = float((pred == y).mean())
        if acc > best["acc"]:
            best = {"acc": acc, "target": float(target)}
            thresholds = thr
    return {str(g): t for g, t in thresholds.items()}


def _nonlinear_recovery(
    df: pd.DataFrame, groups: np.ndarray, protected: str
) -> dict[str, float | str]:
    """Protected-attribute recovery from the features (logistic vs MLP probe).

    Returns {metric, chance, logistic, mlp, note}. Binary -> held-out ROC AUC,
    chance 0.5. Multi-class -> held-out BALANCED accuracy (the mean of the
    per-group recall rates, so every group counts equally whatever its share of
    the rows), chance = 1/K over the K groups present in the whole frame. The
    logistic probe is class-weighted so it cannot collapse to the majority guess.
    A big MLP-over-logistic margin means non-linear recoverability a correlation
    check would miss; both scores should fall toward chance from the biased to
    the debiased features.

    Why not raw accuracy against the majority share: on an uneven characteristic
    an unweighted probe learns to name the largest group for every row and scores
    exactly the majority share before AND after debiasing, so it cannot see a
    change - it reports that the debiasing did nothing. Balanced accuracy scores
    that collapsed probe at chance (1/K), so a real change is visible. The MLP
    has no class-weight option in scikit-learn; it is scored on the same metric
    and can still under-recover a rare group, which the returned note says.
    """
    from sklearn.model_selection import train_test_split

    num, cat = _feature_cols(df, protected)
    pre = ColumnTransformer(
        [("num", StandardScaler(), num), ("cat", OneHotEncoder(handle_unknown="ignore"), cat)]
    )
    x = pre.fit_transform(df[num + cat])
    if hasattr(x, "toarray"):
        x = x.toarray()
    y = LabelEncoder().fit_transform(groups)
    n_groups = int(len(np.unique(y)))  # the frozen full class set, not the test split's
    binary = n_groups == 2
    strat = y if np.bincount(y).min() >= 2 else None
    x_tr, x_te, y_tr, y_te = train_test_split(x, y, test_size=0.3, random_state=1, stratify=strat)
    lin = LogisticRegression(max_iter=2000, class_weight=None if binary else "balanced").fit(
        x_tr, y_tr
    )
    with warnings.catch_warnings():
        # the MLP is a recovery probe, not a tuned model; its optimiser hitting the
        # iteration cap does not affect the biased-vs-debiased comparison.
        warnings.simplefilter("ignore", category=ConvergenceWarning)
        mlp = MLPClassifier(hidden_layer_sizes=(48, 24), max_iter=800, random_state=1).fit(
            x_tr, y_tr
        )
    if binary:
        return {
            "metric": "AUC",
            "chance": 0.5,
            "logistic": float(roc_auc_score(y_te, lin.predict_proba(x_te)[:, 1])),
            "mlp": float(roc_auc_score(y_te, mlp.predict_proba(x_te)[:, 1])),
            "note": "",
        }
    return {
        "metric": "balanced accuracy",
        "chance": 1.0 / n_groups,
        "logistic": float(balanced_accuracy_score(y_te, lin.predict(x_te))),
        "mlp": float(balanced_accuracy_score(y_te, mlp.predict(x_te))),
        "note": (
            "logistic is class-weighted; the MLP is not and can still under-recover a rare group"
        ),
    }


def _fairness_table(pred: np.ndarray, fair_y: np.ndarray, groups: np.ndarray) -> str:
    """Per-group diagnostic table scored against the fair ground truth (any K).

    For each protected group: N, selection rate (predicted-positive rate), TPR,
    FPR, precision, balanced accuracy - the COMPAS-level breakdown behind the
    single group-bias headline. Returns a formatted multi-line string.
    """
    pred = np.asarray(pred).astype(int)
    fair_y = np.asarray(fair_y).astype(int)
    groups = np.asarray(groups)
    lines = [
        f"  {'group':<10} {'N':>6} {'sel':>6} {'TPR':>6} {'FPR':>6} {'prec':>6} {'bal-acc':>8}"
    ]
    for g in sorted(np.unique(groups).tolist()):
        m = groups == g
        p, t = pred[m], fair_y[m]
        pos, neg = t == 1, t == 0
        tpr = float(p[pos].mean()) if pos.any() else float("nan")
        fpr = float(p[neg].mean()) if neg.any() else float("nan")
        tnr = 1.0 - fpr if neg.any() else float("nan")
        prec = float(t[p == 1].mean()) if (p == 1).any() else float("nan")
        bal = (tpr + tnr) / 2 if pos.any() and neg.any() else float("nan")
        lines.append(
            f"  {str(g):<10} {m.sum():>6} {p.mean():>6.3f} {tpr:>6.3f} {fpr:>6.3f} "
            f"{prec:>6.3f} {bal:>8.3f}"
        )
    return "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser(description="Independently confirm Rosa on a synthetic dataset.")
    ap.add_argument("--biased-train", required=True, type=Path)
    ap.add_argument("--raw-holdout", required=True, type=Path)
    ap.add_argument("--debiased-train", default=None, type=Path)
    ap.add_argument("--debiased-holdout", default=None, type=Path)
    ap.add_argument("--protected", default="protected")
    ap.add_argument(
        "--training-manifest",
        default=None,
        type=Path,
        help="Run Manifest JSON of the TRAINING job (pair with --inference-manifest to prove lineage)",
    )
    ap.add_argument(
        "--inference-manifest",
        default=None,
        type=Path,
        help="Run Manifest JSON of the INFERENCE job; its training_job_id must name the training job",
    )
    ap.add_argument(
        "--enforce-gates",
        action="store_true",
        help="exit non-zero if a prespecified scientific gate fails (default: report only)",
    )
    args = ap.parse_args()

    # Checked up front so a mismatched or unreadable manifest pair fails before
    # any compute, not after a full model fit.
    lineage_status = _verify_lineage(args.training_manifest, args.inference_manifest)

    biased = pd.read_csv(args.biased_train)
    raw_hold = pd.read_csv(args.raw_holdout)
    for name, frame in (("biased-train", biased), ("raw-holdout", raw_hold)):
        missing = [c for c in (TARGET, FAIR_TARGET, args.protected) if c not in frame.columns]
        if missing:
            raise SystemExit(f"{name} is missing required column(s): {missing}")
    fair_y = raw_hold[FAIR_TARGET].to_numpy()
    true_groups = raw_hold[args.protected].to_numpy()  # measurement only, never a feature

    # Printed on BOTH paths - a declined run has no debiased output, and the facts
    # are exactly what makes that decline readable.
    eligible_rows = _print_run_facts(biased, args.protected, args.training_manifest)

    # --- Model A: trained on the biased data, scored on the fair ground truth ---
    model_a, cols_a = _fit(biased, TARGET, args.protected)
    proba_a = model_a.predict_proba(raw_hold[cols_a])[:, 1]
    pred_a = (proba_a >= 0.5).astype(int)
    r2_a, brier_a = float(r2_score(fair_y, proba_a)), float(brier_score_loss(fair_y, proba_a))
    acc_a, bias_a = float(accuracy_score(fair_y, pred_a)), group_bias(pred_a, true_groups)

    # --- Model C: per-group parity thresholds on Model A (post-processing) ---
    # proba_c is a CONSTRUCTED score - Model A's probability shifted per group by
    # (0.5 - threshold) so the quota baseline has a comparable continuous output to
    # score R2 against. It is a deliberately naive baseline, NOT a calibrated
    # probability, so NO Brier score is reported for C (Brier needs a real proba).
    thr = _parity(model_a, biased, cols_a, biased[args.protected].to_numpy())
    pred_c = np.array(
        [proba_a[i] >= thr[str(true_groups[i])] for i in range(len(proba_a))], dtype=int
    )
    proba_c = np.clip(
        np.array([proba_a[i] + (0.5 - thr[str(true_groups[i])]) for i in range(len(proba_a))]),
        0.0,
        1.0,
    )
    r2_c = float(r2_score(fair_y, proba_c))
    acc_c, bias_c = float(accuracy_score(fair_y, pred_c)), group_bias(pred_c, true_groups)

    print(f"rows: train={len(biased)} held-out={len(raw_hold)}")
    print("\nprobability metrics (scored on fair_outcome; lower Brier = better calibration):")
    print(f"  A (biased)  R2 {r2_a:+.3f}   Brier {brier_a:.3f}")
    print(f"  C (parity)  R2 {r2_c:+.3f}   Brier   n/a  (constructed score, not calibrated)")
    print("threshold metrics (0.5 cutoff; group-bias = max-min predicted-positive-rate gap):")
    print(f"  A (biased)  accuracy {acc_a:.3f}   group-bias {bias_a:.3f}")
    print(f"  C (parity)  accuracy {acc_c:.3f}   group-bias {bias_c:.3f}")
    print("per-group fairness, Model A (biased), scored on fair_outcome:")
    print(_fairness_table(pred_a, fair_y, true_groups))

    rec_bi = _nonlinear_recovery(raw_hold, true_groups, args.protected)
    print(
        f"protected recovery ({rec_bi['metric']}, chance={rec_bi['chance']:.3f})  "
        f"biased features:   logistic={rec_bi['logistic']:.3f} MLP={rec_bi['mlp']:.3f}"
    )
    if rec_bi["note"]:
        print(f"  ({rec_bi['note']})")

    if args.debiased_train is None or args.debiased_holdout is None:
        print("\nModel B needs Rosa's debiased output - run a Training + Inference job, then")
        print("re-run with --debiased-train / --debiased-holdout.")
        if eligible_rows < UTILITY_CLAIM_MIN_TRAINING_ROWS:
            print()
            print(UTILITY_FLOOR_NOTE)
        return

    deb_train = pd.read_csv(args.debiased_train)
    deb_hold = pd.read_csv(args.debiased_holdout)
    # Structural gates, ALL fail closed, before any Model B metric or by-index
    # group re-attach. They run in order of what they prove:
    #   1. ALIGNMENT - the unique per-row audit_row_id rides through Rosa as an
    #      ignored pass-through, so a row-for-row match proves the order held.
    #   2. INTEGRITY - every other ignored column also survived unchanged, so a
    #      flipped outcome cannot silently poison the scoring key.
    #   3. PROVENANCE-BY-SHAPE - the protected column really is gone.
    #   4. LINEAGE - and, if the manifests were supplied, that these two outputs
    #      came from the SAME Rosa run rather than merely the same input rows.
    _explain_sampled_training(biased, deb_train, args.training_manifest)
    _verify_row_ids(biased, deb_train, "training output")
    _verify_row_ids(raw_hold, deb_hold, "holdout output")
    passthrough = [c for c in (TARGET, FAIR_TARGET, ROW_ID) if c in biased.columns]
    _verify_passthrough(biased, deb_train, passthrough, "training output")
    _verify_passthrough(raw_hold, deb_hold, passthrough, "holdout output")
    _verify_protected_absent(deb_train, args.protected, "training output")
    _verify_protected_absent(deb_hold, args.protected, "holdout output")
    print(lineage_status)
    missing_cols = [c for c in cols_a if c not in deb_train.columns or c not in deb_hold.columns]
    if missing_cols:  # model columns are selected by NAME; the debiased frames must carry them
        raise SystemExit(f"debiased output is missing model feature column(s): {missing_cols}")

    model_b, cols_b = _fit(deb_train, TARGET, args.protected)
    proba_b = model_b.predict_proba(deb_hold[cols_b])[:, 1]
    pred_b = (proba_b >= 0.5).astype(int)
    r2_b, brier_b = float(r2_score(fair_y, proba_b)), float(brier_score_loss(fair_y, proba_b))
    acc_b, bias_b = float(accuracy_score(fair_y, pred_b)), group_bias(pred_b, true_groups)
    print(
        f"\nB (Rosa)    probability: R2 {r2_b:+.3f}  Brier {brier_b:.3f}   "
        f"[stochastic; run 10 independent Rosa runs, take the mean]"
    )
    print(f"B (Rosa)    threshold:   accuracy {acc_b:.3f}  group-bias {bias_b:.3f}")
    print("per-group fairness, Model B (Rosa), scored on fair_outcome:")
    print(_fairness_table(pred_b, fair_y, true_groups))

    deb_hold_grp = deb_hold.copy()
    deb_hold_grp[args.protected] = true_groups  # re-attach by index (verified aligned above)
    rec_db = _nonlinear_recovery(deb_hold_grp, true_groups, args.protected)
    print(
        f"protected recovery ({rec_db['metric']}, chance={rec_db['chance']:.3f})  "
        f"debiased features: logistic={rec_db['logistic']:.3f} MLP={rec_db['mlp']:.3f}  "
        "(should fall toward chance)"
    )
    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)")

    # --- prespecified scientific gates -------------------------------------
    # Reported by default; --enforce-gates makes a failure exit non-zero. The
    # split is deliberate: a failed ordering on your own data can be a genuine
    # result, but a script used as an approval gate has to be able to say NO.
    # The utility gate is THREE-STATE. Below the measured floor the ordering is
    # close to a coin flip, so this script renders NO VERDICT rather than a pass or
    # a fail - and it never infers a pass from the row count alone, in either
    # direction. A row count above the floor is a precondition for reading the
    # result, not the result.
    below_floor = eligible_rows < UTILITY_CLAIM_MIN_TRAINING_ROWS
    ordering_ok = ordering == "C < A < B"
    multi_class = int(biased[args.protected].nunique()) > 2

    print("\nprespecified gates:")
    if below_floor:
        print(
            f"  [NO VERDICT] ordering C < A < B by R2  (got {ordering}; "
            f"{eligible_rows} eligible training rows is below the characterised "
            f"range of {UTILITY_CLAIM_MIN_TRAINING_ROWS}+)"
        )
        print("               outside the characterised utility range - see the note below.")
    else:
        print(f"  [{'PASS' if ordering_ok else 'FAIL'}] ordering C < A < B by R2  (got {ordering})")
    if multi_class:
        print(MULTICLASS_SCOPE_NOTE)

    recovery_ok = rec_db["mlp"] < rec_bi["mlp"]
    print(
        f"  [{'PASS' if recovery_ok else 'FAIL'}] protected recoverability falls after "
        f"debiasing  (MLP recovery {rec_bi['mlp']:.3f} -> {rec_db['mlp']:.3f})"
    )

    print(
        f"\nThis is ONE Rosa run. The published protocol is {PUBLISHED_PROTOCOL_RUNS} independent\n"
        "Rosa runs, aggregated as a mean - Rosa exposes no training seed, so the runs ARE the\n"
        "measurement and a single run cannot settle the ordering on its own."
    )
    if below_floor:
        print()
        print(UTILITY_FLOOR_NOTE)

    outcomes = (
        ("ordering C < A < B by R2", ordering_ok and not below_floor),
        ("protected recoverability falls after debiasing", recovery_ok),
    )
    failed = [name for name, ok in outcomes if not ok]
    if failed:
        if args.enforce_gates:
            # A NO VERDICT is not a pass: a script used as an approval gate must
            # decline to approve when it cannot render a verdict, exactly as it
            # declines when a gate fails.
            reason = (
                "below the characterised utility range"
                if below_floor
                else f"{len(failed)} prespecified gate(s) failed"
            )
            raise SystemExit(f"\n{reason}: {', '.join(failed)}")
        print(
            f"\n{len(failed)} gate(s) did not pass. This exits 0 because a failed or "
            "unrenderable gate can be a real result on your data; re-run with "
            "--enforce-gates to make it exit non-zero."
        )


if __name__ == "__main__":
    main()
