import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.metrics import average_precision_score, precision_recall_curve, roc_auc_score, roc_curve

# The SAME scorer at three prevalences: healthy scores ~ N(0, 1), sick scores ~ N(2, 1).
# Only the number of sick patients changes (same code as in the lesson).
rng = np.random.default_rng(0)
neg = rng.normal(0, 1, 100_000)
t = 1.16                                     # one fixed threshold, marked on every curve
label_at = {0.5: (0.48, 0.91), 0.1: (0.36, 0.47), 0.01: (0.3, 0.035)}
colors = {0.5: '#58a6ff', 0.1: '#f0883e', 0.01: '#ff7b72'}

fig, (left, right) = plt.subplots(1, 2, figsize=(9.5, 5.0))
grid = np.linspace(0, 1, 400)
for prev, c in colors.items():
    n_pos = round(len(neg) * prev / (1 - prev))
    pos = rng.normal(2, 1, n_pos)
    s, y = np.r_[pos, neg], np.r_[np.ones(n_pos), np.zeros(len(neg))]
    fpr, tpr, _ = roc_curve(y, s)
    prec, rec, _ = precision_recall_curve(y, s)
    auc, ap = roc_auc_score(y, s), average_precision_score(y, s)
    left.plot(grid, np.interp(grid, fpr, tpr), color=c, lw=1.8,
              label=f'{prev:.0%} sick ({n_pos:,} of {n_pos + len(neg):,}): AUC = {auc:.3f}')
    right.plot(grid, np.interp(grid, rec[::-1], prec[::-1]), color=c, lw=1.8,
               label=f'{prev:.0%} sick: AP = {ap:.3f}')
    right.axhline(prev, color=c, ls=':', lw=1)
    tp, fp = (pos >= t).sum(), (neg >= t).sum()
    left.plot(fp / len(neg), tp / n_pos, 'o', color=c, ms=7, mfc='none', mew=1.8)
    right.plot(tp / n_pos, tp / (tp + fp), 'o', color=c, ms=7, mfc='none', mew=1.8)
    right.annotate(f'precision {tp / (tp + fp):.2f}', xy=(tp / n_pos, tp / (tp + fp)),
                   xytext=label_at[prev], fontsize=8, color=c,
                   arrowprops=dict(arrowstyle='->', color=c, lw=0.8))
left.plot([0, 1], [0, 1], '--', color='gray', lw=1)
left.annotate(f'threshold {t}: TPR ≈ 0.80, FPR ≈ 0.12\n(the same point for all three)', xy=(0.125, 0.8),
              xytext=(0.42, 0.15), fontsize=8, color='gray', arrowprops=dict(arrowstyle='->', color='gray'))
left.set_xlabel('false positive rate  (FP / all healthy)')
left.set_ylabel('true positive rate (recall)')
left.set_title('ROC: the three curves coincide', fontsize=10)
left.legend(fontsize=7.5, loc='upper center', bbox_to_anchor=(0.5, -0.14), ncol=1)
right.text(0.99, 0.5 + 0.015, 'dotted: random scorer = prevalence', fontsize=7.5, color='gray', ha='right')
right.set_xlabel('recall')
right.set_ylabel('precision  (TP / all flagged)')
right.set_title('Precision–recall: the curve sinks with the prevalence', fontsize=10)
right.legend(fontsize=7.5, loc='upper center', bbox_to_anchor=(0.5, -0.14), ncol=3, handlelength=1.2,
             columnspacing=1)
for ax in (left, right):
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1.02)
    ax.grid(color='gray', alpha=0.2, lw=0.5)

plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()