#!/usr/bin/env python3
"""Independent audit of fair decisions from calibrated finite scores.

The paper's released boundary tracer is run at frozen commit 9d8b29f and
audited against a separately formulated constrained optimizer over every
per-bin randomized decision probability.
"""

from __future__ import annotations

import itertools
import json
import time
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize

from boundary_trace import GroupScoreDistribution, trace_intersection


HERE = Path(__file__).resolve().parent
OUTPUTS = HERE / "outputs" / "fair_calibrated"
OUTPUTS.mkdir(parents=True, exist_ok=True)
SEED = 260207285


def group_stats(rule: np.ndarray, dist: GroupScoreDistribution) -> dict:
    score = np.asarray(dist.s_input, dtype=float)
    weight = np.asarray(dist.w_input, dtype=float)
    select = np.asarray(rule, dtype=float)
    tp = float(np.sum(weight * score * select))
    fp = float(np.sum(weight * (1.0 - score) * select))
    fn = float(np.sum(weight * score * (1.0 - select)))
    tn = float(np.sum(weight * (1.0 - score) * (1.0 - select)))
    return {
        "tp": tp,
        "fp": fp,
        "fn": fn,
        "tn": tn,
        "ppv": tp / (tp + fp),
        "for": fn / (fn + tn),
        "tpr": tp / (tp + fn),
        "fpr": fp / (fp + tn),
        "accuracy": tp + tn,
    }


def aggregate(rules: list[np.ndarray], dists: list[GroupScoreDistribution], prob_a1: float) -> dict:
    stats = [group_stats(rule, dist) for rule, dist in zip(rules, dists)]
    group_weights = np.array([1.0 - prob_a1, prob_a1])
    accuracy = float(sum(group_weights[a] * stats[a]["accuracy"] for a in range(2)))
    prevalences = np.array([dist.pi for dist in dists])
    prevalence = float(group_weights @ prevalences)
    pooled_tpr = float(sum(group_weights[a] * prevalences[a] * stats[a]["tpr"] for a in range(2)) / prevalence)
    pooled_fpr = float(sum(group_weights[a] * (1.0 - prevalences[a]) * stats[a]["fpr"] for a in range(2)) / (1.0 - prevalence))
    deviation = float(
        sum(
            group_weights[a]
            * (
                prevalences[a] * abs(stats[a]["tpr"] - pooled_tpr)
                + (1.0 - prevalences[a]) * abs(stats[a]["fpr"] - pooled_fpr)
            )
            for a in range(2)
        )
    )
    return {
        "group_stats": stats,
        "accuracy": accuracy,
        "separation_tv": deviation,
        "ppv_gap": abs(stats[0]["ppv"] - stats[1]["ppv"]),
        "for_gap": abs(stats[0]["for"] - stats[1]["for"]),
    }


def parity_cross_products(x: np.ndarray, dists: list[GroupScoreDistribution]) -> np.ndarray:
    cut = len(dists[0].s_input)
    a = group_stats(x[:cut], dists[0])
    b = group_stats(x[cut:], dists[1])
    ppv = a["tp"] * (b["tp"] + b["fp"]) - b["tp"] * (a["tp"] + a["fp"])
    omission = a["fn"] * (b["fn"] + b["tn"]) - b["fn"] * (a["fn"] + a["tn"])
    return np.array([ppv, omission])


def independent_optimum(
    dists: list[GroupScoreDistribution],
    prob_a1: float,
    objective: str,
    official_rules: list[np.ndarray],
    rng: np.random.Generator,
) -> dict:
    sizes = [len(d.s_input) for d in dists]
    nvar = sum(sizes)
    clipped_official = np.clip(np.concatenate(official_rules), 1e-9, 1.0 - 1e-9)
    starts = [clipped_official]
    # These starts contain no boundary-tracer information.  They make the
    # comparison an optimizer audit rather than merely reevaluating its output.
    starts += [rng.uniform(0.02, 0.98, nvar) for _ in range(13)]
    constraint = {"type": "eq", "fun": lambda z: parity_cross_products(z, dists)}
    feasible = []

    def unpack(z: np.ndarray) -> list[np.ndarray]:
        return [z[: sizes[0]], z[sizes[0] :]]

    for start in starts:
        def loss(z: np.ndarray) -> float:
            result = aggregate(unpack(z), dists, prob_a1)
            return -result["accuracy"] if objective == "accuracy" else result["separation_tv"]

        out = minimize(
            loss,
            start,
            method="SLSQP",
            bounds=[(1e-9, 1.0 - 1e-9)] * nvar,
            constraints=constraint,
            options={"ftol": 2e-12, "maxiter": 1200},
        )
        violation = float(np.max(np.abs(parity_cross_products(out.x, dists))))
        if out.success and violation < 2e-8:
            result = aggregate(unpack(out.x), dists, prob_a1)
            feasible.append((out.fun, result, violation, out.x))
    if not feasible:
        raise AssertionError("independent constrained optimizer found no feasible solution")
    best = min(feasible, key=lambda row: row[0])
    return {
        "feasible_multistarts": len(feasible),
        "objective": objective,
        "accuracy": best[1]["accuracy"],
        "separation_tv": best[1]["separation_tv"],
        "ppv_gap": best[1]["ppv_gap"],
        "for_gap": best[1]["for_gap"],
        "max_crossproduct_violation": best[2],
        "rule": best[3].tolist(),
    }


def paper_cases() -> list[tuple[str, list[GroupScoreDistribution], float]]:
    shared = GroupScoreDistribution([0.1, 0.2, 0.5, 0.7, 0.9], [0.1, 0.3, 0.3, 0.15, 0.15], name="paper group 0")
    group_a = GroupScoreDistribution([0.12, 0.3, 0.85], [0.15, 0.45, 0.4], name="paper group 1a")
    group_b = GroupScoreDistribution([0.2, 0.3, 0.75], [0.3, 0.5, 0.2], name="paper group 1b")
    return [("paper_A", [shared, group_a], 0.5), ("paper_B", [shared, group_b], 0.5)]


def fresh_cases(rng: np.random.Generator, count: int = 4) -> list[tuple[str, list[GroupScoreDistribution], float]]:
    cases = []
    attempts = 0
    while len(cases) < count and attempts < 1000:
        attempts += 1
        dists = []
        for group in range(2):
            bins = int(rng.integers(4, 7))
            scores = np.sort(rng.uniform(0.03, 0.97, bins))
            scores[0] = min(scores[0], 0.12 + 0.02 * group)
            scores[-1] = max(scores[-1], 0.88 - 0.02 * group)
            weights = rng.dirichlet(np.full(bins, 1.7))
            dists.append(GroupScoreDistribution(scores, weights, name=f"fresh {len(cases)} group {group}"))
        prob_a1 = float(rng.uniform(0.3, 0.7))
        try:
            trace_intersection(dists[0], dists[1], prob_a1)
        except (ValueError, IndexError):
            continue
        if abs(dists[0].pi - dists[1].pi) < 0.035:
            continue
        cases.append((f"fresh_{len(cases) + 1}", dists, prob_a1))
    if len(cases) != count:
        raise AssertionError("could not construct enough nontrivial fresh intersections")
    return cases


def deterministic_control(dists: list[GroupScoreDistribution], prob_a1: float) -> dict:
    gaps = []
    checked = 0
    for left in itertools.product((0.0, 1.0), repeat=len(dists[0].s_input)):
        if sum(left) in (0, len(left)):
            continue
        for right in itertools.product((0.0, 1.0), repeat=len(dists[1].s_input)):
            if sum(right) in (0, len(right)):
                continue
            out = aggregate([np.array(left), np.array(right)], dists, prob_a1)
            gaps.append(max(out["ppv_gap"], out["for_gap"]))
            checked += 1
    return {
        "nonconstant_deterministic_rule_pairs": checked,
        "minimum_max_predictive_parity_gap": min(gaps),
    }


def audit(rng: np.random.Generator) -> dict:
    cases = paper_cases() + fresh_cases(rng)
    rows = []
    max_calibration_residual = 0.0
    max_official_parity_gap = 0.0
    max_accuracy_gap = 0.0
    max_separation_gap = 0.0
    threshold_counterexamples = 0
    max_threshold_gap = 0.0

    for name, dists, prob_a1 in cases:
        for dist in dists:
            # In the exact joint construction P(Y=1,S=s)/P(S=s)=s.
            joint_positive = dist.w_input * dist.s_input
            recovered_score = joint_positive / dist.w_input
            max_calibration_residual = max(max_calibration_residual, float(np.max(np.abs(recovered_score - dist.s_input))))

        threshold_rules = [(dist.s_input >= 0.5).astype(float) for dist in dists]
        threshold = aggregate(threshold_rules, dists, prob_a1)
        threshold_gap = max(threshold["ppv_gap"], threshold["for_gap"])
        threshold_counterexamples += threshold_gap > 1e-3
        max_threshold_gap = max(max_threshold_gap, threshold_gap)

        official = trace_intersection(dists[0], dists[1], prob_a1)
        case_row = {
            "case": name,
            "group_prevalences": [float(d.pi) for d in dists],
            "group_score_bins": [len(d.s_input) for d in dists],
            "probability_group_1": prob_a1,
            "common_threshold": {
                "accuracy": threshold["accuracy"],
                "ppv_gap": threshold["ppv_gap"],
                "for_gap": threshold["for_gap"],
            },
        }
        for objective, optimum in (("accuracy", official.max_acc), ("separation", official.min_dsep)):
            official_rules = [dist.selection_rule(optimum.p, optimum.q) for dist in dists]
            checked = aggregate(official_rules, dists, prob_a1)
            independent = independent_optimum(dists, prob_a1, objective, official_rules, rng)
            official_value = checked["accuracy"] if objective == "accuracy" else checked["separation_tv"]
            independent_value = independent["accuracy"] if objective == "accuracy" else independent["separation_tv"]
            gap = abs(official_value - independent_value)
            if objective == "accuracy":
                max_accuracy_gap = max(max_accuracy_gap, gap)
            else:
                max_separation_gap = max(max_separation_gap, gap)
            max_official_parity_gap = max(max_official_parity_gap, checked["ppv_gap"], checked["for_gap"])
            case_row[objective] = {
                "common_ppv_p": float(optimum.p),
                "common_for_q": float(optimum.q),
                "accuracy": checked["accuracy"],
                "separation_tv": checked["separation_tv"],
                "ppv_gap": checked["ppv_gap"],
                "for_gap": checked["for_gap"],
                "independent_optimizer": independent,
                "objective_gap_vs_independent": gap,
                "randomized_rules": [rule.tolist() for rule in official_rules],
            }
        rows.append(case_row)

    control = deterministic_control(cases[0][1], cases[0][2])
    assert max_calibration_residual < 1e-15
    assert threshold_counterexamples == len(cases)
    assert max_official_parity_gap < 2e-12
    assert max_accuracy_gap < 2e-7
    assert max_separation_gap < 2e-7
    assert control["minimum_max_predictive_parity_gap"] > 1e-3
    return {
        "cases": rows,
        "summary": {
            "total_cases": len(cases),
            "fresh_cases": len(cases) - 2,
            "max_exact_calibration_residual": max_calibration_residual,
            "common_threshold_counterexamples": int(threshold_counterexamples),
            "maximum_threshold_predictive_parity_gap": max_threshold_gap,
            "max_sufficient_rule_ppv_or_for_gap": max_official_parity_gap,
            "max_accuracy_objective_gap_vs_independent": max_accuracy_gap,
            "max_separation_objective_gap_vs_independent": max_separation_gap,
        },
        "negative_control": control,
    }


def make_figure(summary: dict, path: Path) -> None:
    rows = summary["claims"]["cases"]
    names = [row["case"] for row in rows]
    threshold = [max(row["common_threshold"]["ppv_gap"], row["common_threshold"]["for_gap"]) for row in rows]
    fair = [max(row["accuracy"]["ppv_gap"], row["accuracy"]["for_gap"]) for row in rows]
    acc = [row["accuracy"]["accuracy"] for row in rows]
    sep = [row["separation"]["separation_tv"] for row in rows]
    x = np.arange(len(rows))
    fig, axes = plt.subplots(1, 3, figsize=(14.3, 4.2))
    axes[0].bar(x - 0.18, threshold, 0.36, label="threshold", color="#dc2626")
    axes[0].bar(x + 0.18, fair, 0.36, label="sufficient optimum", color="#0d9488")
    axes[0].set_yscale("log")
    axes[0].set_xticks(x, names, rotation=35, ha="right")
    axes[0].set_ylabel("max(PPV gap, FOR gap)")
    axes[0].set_title("Calibration does not survive thresholding")
    axes[0].legend(frameon=False)
    axes[0].grid(axis="y", alpha=0.25)

    axes[1].bar(x, acc, color="#4f46e5")
    axes[1].set_ylim(0, 1)
    axes[1].set_xticks(x, names, rotation=35, ha="right")
    axes[1].set_ylabel("accuracy")
    axes[1].set_title("Exact sufficient optimum")
    axes[1].grid(axis="y", alpha=0.25)

    axes[2].bar(x, sep, color="#f59e0b")
    axes[2].set_xticks(x, names, rotation=35, ha="right")
    axes[2].set_ylabel("TV deviation from separation")
    axes[2].set_title("Minimum separation deviation")
    axes[2].grid(axis="y", alpha=0.25)

    fig.suptitle("Fair decisions from calibrated scores — exact boundary audit", fontsize=14, fontweight="bold")
    fig.tight_layout()
    fig.savefig(path, dpi=180, bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    rng = np.random.default_rng(SEED)
    start = time.perf_counter()
    summary = {
        "paper": {
            "title": "Fair Decisions from Calibrated Scores",
            "openreview_id": "XiD5RhcEDK",
            "arxiv_id": "2602.07285",
        },
        "audit": {
            "seed": SEED,
            "official_boundary_trace_commit": "9d8b29f116f627868e9d5f2a31baa6d8567a0920",
            "independent_check": "multistart constrained optimization over every randomized bin decision",
        },
        "claims": audit(rng),
    }
    summary["audit"]["wall_seconds"] = time.perf_counter() - start
    (OUTPUTS / "fair_calibrated_results.json").write_text(json.dumps(summary, indent=2) + "\n")
    make_figure(summary, OUTPUTS / "fair_calibrated_audit.png")
    print(json.dumps(summary["claims"]["summary"], indent=2))
    print(json.dumps(summary["claims"]["negative_control"], indent=2))


if __name__ == "__main__":
    main()