import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.metrics import average_precision_score

# The same 20 patients, sorted by score
score = np.array([0.95, 0.90, 0.85, 0.80, 0.75, 0.70, 0.65, 0.60, 0.55, 0.45,
                  0.40, 0.35, 0.30, 0.25, 0.20, 0.15, 0.10, 0.08, 0.05, 0.02])
sick = np.array([1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0])
tp = np.cumsum(sick)                      # threshold = score of patient k: k patients flagged
flagged = np.arange(1, 21)
precision, recall = tp / flagged, tp / sick.sum()
ap = average_precision_score(sick, score)

fig, ax = plt.subplots(figsize=(7.5, 4.4))
# average precision: one rectangle of width 1/10 per sick patient, height = precision there
for k in np.where(sick == 1)[0]:
    ax.add_patch(plt.Rectangle((recall[k] - 0.1, 0), 0.1, precision[k], color='#58a6ff', alpha=0.13, lw=0))
    ax.plot([recall[k] - 0.1, recall[k]], [precision[k]] * 2, color='#58a6ff', lw=1, alpha=0.6)
ax.plot(recall, precision, '-', color='gray', lw=1, alpha=0.7, zorder=2)
ax.scatter(recall[sick == 1], precision[sick == 1], color='#3fb950', s=26, zorder=3,
           label='after a sick patient: recall +0.1, precision rises')
ax.scatter(recall[sick == 0], precision[sick == 0], color='#ff7b72', s=26, zorder=3,
           label='after a healthy patient: recall unchanged, precision drops')
for k in [4, 7, 10, 12, 13]:
    ax.text(recall[k] + 0.012, precision[k], f'P{k + 1:02d}', fontsize=7.5, color='#ff7b72', va='center')
ax.text(recall[19] + 0.012, 0.565, 'P16…P20', fontsize=7.5, color='#ff7b72', va='center')
ax.axhline(0.5, color='gray', ls='--', lw=1)
ax.text(0.02, 0.47, 'random scorer: precision = prevalence = 10/20 = 0.5', fontsize=8, color='gray', va='top')
k = 8                                         # threshold 0.5: the first 9 patients flagged
ax.plot(recall[k], precision[k], 'o', ms=11, mfc='none', mec='#f0883e', mew=2)
ax.annotate(f'threshold 0.5: recall 0.7,\nprecision 7/9 = {precision[k]:.3f}', xy=(recall[k], precision[k]),
            xytext=(0.43, 0.6), fontsize=8.5, color='#f0883e', arrowprops=dict(arrowstyle='->', color='#f0883e'))
ax.text(0.03, 0.22, f'shaded area = average precision\n= (1+1+1+1+5/6+6/7+7/9+8/10+9/12+10/15) / 10\n= {ap:.4f}',
        fontsize=8, color='#58a6ff')
ax.set_xlim(0, 1.1)
ax.set_ylim(0, 1.05)
ax.set_xticks(np.arange(0, 1.01, 0.1))
ax.grid(color='gray', alpha=0.2, lw=0.5)
ax.set_xlabel('recall = TP / 10 sick')
ax.set_ylabel('precision = TP / flagged')
ax.legend(fontsize=7.5, loc='lower right')
ax.set_title('Precision–recall of the same 20 patients: a saw-tooth, not a staircase', fontsize=9.5)

plt.tight_layout()

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

plt.close()