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})

# The same test (sensitivity 0.70, specificity 0.80: our model at threshold 0.5)
# applied to populations where the disease is more or less common
sens, spec = 0.70, 0.80
prev = np.logspace(-3, np.log10(0.5), 400)
precision = sens * prev / (sens * prev + (1 - spec) * (1 - prev))
npv = spec * (1 - prev) / (spec * (1 - prev) + (1 - sens) * prev)
accuracy = sens * prev + spec * (1 - prev)

fig, ax = plt.subplots(figsize=(8.4, 4.0))
ax.plot(prev, np.full_like(prev, sens), color=ORANGE, ls='--', lw=1.4, label='recall (sensitivity) = 0.70')
ax.plot(prev, np.full_like(prev, spec), color=GRAY, ls='--', lw=1.4, label='specificity = 0.80')
ax.plot(prev, accuracy, color=RED, lw=1.4, label='accuracy')
ax.plot(prev, npv, color=GREEN, lw=2, label='NPV')
ax.plot(prev, precision, color=BLUE, lw=2.4, label='precision')
for p0, dx, va in [(0.5, -0.03, 'bottom'), (0.1, 0.0, 'bottom'), (0.01, 0.0, 'bottom')]:
    pr = sens * p0 / (sens * p0 + (1 - spec) * (1 - p0))
    ax.scatter([p0], [pr], color=BLUE, s=40, zorder=5)
    ax.text(p0 * 1.18, pr + (0.03 if p0 < 0.5 else -0.03), f'{pr:.3f}', color=BLUE, fontsize=9,
            ha='left', va=va if p0 < 0.5 else 'top', weight='bold')
    ax.axvline(p0, color=GRAY, ls=':', lw=0.8)
ax.set_xscale('log')
ax.set_xticks([0.001, 0.01, 0.1, 0.5], ['0.1%', '1%', '10%', '50%'])
ax.set_xlim(0.001, 0.8)
ax.set_ylim(0, 1.05)
ax.set_xlabel('prevalence: share of the population that is sick (log scale)')
ax.set_ylabel('metric')
ax.legend(fontsize=8, loc='center left', frameon=False, bbox_to_anchor=(1.0, 0.5))
ax.set_title('Same test, rarer disease: precision collapses, recall and specificity do not move',
             fontsize=9.5)
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()