import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

BLUE, ORANGE, GREEN, RED, GRAY = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72', '#8b949e'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})

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])
y_true = np.array([1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0])


def prec_rec(t):
    pred = score >= t
    tp = (pred & (y_true == 1)).sum()
    flagged = pred.sum()
    return (tp / flagged if flagged else np.nan), tp / (y_true == 1).sum()


# precision and recall only change when the threshold crosses a patient's score
grid = np.linspace(0.0, 1.0, 2001)
P, R = np.array([prec_rec(t) for t in grid]).T

fig, ax = plt.subplots(figsize=(8.6, 3.9))
ax.plot(grid, R, color=ORANGE, lw=2, label='recall  (of the 10 sick, how many are flagged)')
ax.plot(grid, P, color=BLUE, lw=2, label='precision  (of the flagged, how many are sick)')
lines = []
for t in [0.3, 0.5, 0.7]:
    p, r = prec_rec(t)
    ax.vlines(t, 0.42, 1.03, color=GRAY, ls=':', lw=1)
    ax.scatter([t], [p], color=BLUE, zorder=5, s=30)
    ax.scatter([t], [r], color=ORANGE, zorder=5, s=30)
    ax.text(t, 1.055, f't = {t}', ha='center', fontsize=8)
    lines.append(f't = {t}:  precision {p:.3f}   recall {r:.3f}')
ax.text(0.34, 0.06, '\n'.join(lines), fontsize=8.5, family='monospace', va='bottom')
# a rug of the 20 patients: each step in a curve is one patient crossing the line
for s, t in zip(score, y_true):
    ax.plot([s, s], [-0.12, -0.04], color=GREEN if t else GRAY, lw=2)
ax.annotate('low threshold: flag almost everyone\n→ recall 1.0, precision falls to 0.5',
            xy=(0.03, 0.52), xytext=(0.03, 0.2), fontsize=8, arrowprops=dict(arrowstyle='->', color=GRAY))
ax.annotate('high threshold: flag only\nthe surest → precision 1.0,\nrecall falls',
            xy=(0.86, 0.99), xytext=(0.785, 0.5), fontsize=8, arrowprops=dict(arrowstyle='->', color=GRAY))
ax.set_xlim(0, 1)
ax.set_ylim(-0.15, 1.12)
ax.set_yticks(list(np.arange(0, 1.01, 0.2)) + [-0.08],
              [f'{v:.1f}' for v in np.arange(0, 1.01, 0.2)] + ['patients'])
ax.set_xlabel('threshold t   (flag a patient when score ≥ t)')
ax.legend(fontsize=8, loc='upper center', frameon=False, ncol=2, bbox_to_anchor=(0.5, 1.16))
for side in ['top', 'right']:
    ax.spines[side].set_visible(False)

plt.tight_layout()

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

plt.close()