# 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 numpy as np
from sklearn.metrics import confusion_matrix

def observations(tn, fp, fn, tp):
    truth = np.array([0] * (tn + fp) + [1] * (fn + tp))
    pred = np.array([0] * tn + [1] * fp + [0] * fn + [1] * tp)
    return truth, pred

def safe_ratio(numerator, denominator):
    return numerator / denominator if denominator else np.nan

def report(truth, pred):
    tn, fp, fn, tp = confusion_matrix(truth, pred, labels=[0, 1]).ravel()
    return {"n": len(truth), "positives": int(tp + fn), "selected": int(tp + fp),
        "tpr": safe_ratio(tp, tp + fn), "fpr": safe_ratio(fp, fp + tn),
        "ppv": safe_ratio(tp, tp + fp), "selection": safe_ratio(tp + fp, len(truth))}

groups = {"A": observations(56, 14, 6, 24),
          "B": observations(72, 18, 2, 8), "C": observations(4, 0, 0, 0)}
metrics = {name: report(*data) for name, data in groups.items()}
assert np.isclose(metrics["A"]["tpr"], metrics["B"]["tpr"])
assert np.isclose(metrics["A"]["fpr"], metrics["B"]["fpr"])
assert not np.isclose(metrics["A"]["ppv"], metrics["B"]["ppv"])
assert np.isnan(metrics["C"]["tpr"]) and np.isnan(metrics["C"]["ppv"])
assert metrics["C"]["fpr"] == 0.0
for name, values in metrics.items():
    print(name, values)

def wilson(successes, total, z=1.96):
    if total == 0:
        return (np.nan, np.nan)
    rate = successes / total
    denominator = 1.0 + z * z / total
    center = (rate + z * z / (2 * total)) / denominator
    half = z * np.sqrt(rate * (1 - rate) / total + z * z / (4 * total**2)) / denominator
    return center - half, center + half

small = wilson(8, 10)
large = wilson(80, 100)
assert small[1] - small[0] > large[1] - large[0]
assert 0 <= small[0] < 0.8 < small[1] <= 1
print("Recall 8/10 interval", small, "80/100 interval", large)