# Source: content/notes/ml/ethics-and-fairness.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import itertools
import numpy as np
from sklearn.metrics import confusion_matrix

rng = np.random.default_rng(10)
n = 2400
group = np.tile(np.array([0, 1]), n // 2)
score = np.where(group == 0, rng.beta(3, 3, n), rng.beta(2, 5, n))
y = rng.binomial(1, score)
order = rng.permutation(n)
validation, test = order[:1200], order[1200:]

def evaluate(indices, thresholds):
    prediction = (score[indices] >= np.asarray(thresholds)[group[indices]]).astype(int)
    rates = []
    cost = 0
    for g in [0, 1]:
        mask = group[indices] == g
        tn, fp, fn, tp = confusion_matrix(y[indices][mask], prediction[mask], labels=[0, 1]).ravel()
        assert tp + fn > 0 and fp + tn > 0
        rates.append((tp / (tp + fn), fp / (fp + tn)))
        cost += fp + 2 * fn
    gap = np.max(np.abs(np.asarray(rates[0]) - np.asarray(rates[1])))
    return float(cost / len(indices)), float(gap)

grid = list(itertools.product(np.linspace(0, 1, 11), repeat=2))
records = [(thresholds, *evaluate(validation, thresholds)) for thresholds in grid]
feasible = [row for row in records if row[2] <= 0.10]
assert feasible, "Always-positive and always-negative policies should be feasible here"
chosen = min(feasible, key=lambda row: (row[1], row[2], row[0]))
frontier = [row for row in records if not any(
    other[1] <= row[1] and other[2] <= row[2] and
    (other[1] < row[1] or other[2] < row[2]) for other in records)]
test_cost, test_gap = evaluate(test, chosen[0])
assert chosen[2] <= 0.10
assert len(frontier) > 0 and np.isfinite([test_cost, test_gap]).all()
print("thresholds", chosen[0], "validation cost/gap", chosen[1:])
print("test cost/gap", (test_cost, test_gap), "frontier points", len(frontier))