import matplotlib.pyplot as plt
from io import StringIO

# Three ways to cut the 20 patients. For each: precision and recall (the two ends of a bar),
# their arithmetic mean (always the middle of the bar) and F1 (pulled toward the smaller end).
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})

cuts = [('flag only P01\n(t = 0.95)', 1 / 1, 1 / 10),
        ('threshold 0.5', 7 / 9, 7 / 10),
        ('flag all 20\n(t = 0.02)', 10 / 20, 10 / 10)]

fig, ax = plt.subplots(figsize=(7.6, 4.4))
for row, (name, p, r) in enumerate(cuts):
    y = len(cuts) - 1 - row
    mean, f1 = (p + r) / 2, 2 * p * r / (p + r)
    ax.plot([min(p, r), max(p, r)], [y, y], color=GRAY, lw=5, alpha=0.25, solid_capstyle='round')
    ax.scatter([p], [y], s=80, color=ORANGE, zorder=3, label='precision' if row == 0 else None)
    ax.scatter([r], [y], s=80, color=GREEN, zorder=3, label='recall' if row == 0 else None)
    ax.scatter([mean], [y], s=110, facecolors='none', edgecolors=GRAY, lw=1.8, zorder=4,
               label='arithmetic mean (middle of the bar)' if row == 0 else None)
    ax.scatter([f1], [y], s=120, marker='D', color=BLUE, zorder=5,
               label='F1 (harmonic mean)' if row == 0 else None)
    ax.text(1.04, y + 0.12, f'mean {mean:.3f}', fontsize=10.5, color=GRAY, va='center')
    ax.text(1.04, y - 0.14, f'F1 {f1:.3f}', fontsize=10.5, color=BLUE, va='center', fontweight='bold')
ax.set_yticks(range(len(cuts)))
ax.set_yticklabels([c[0] for c in cuts][::-1], fontsize=10.5)
ax.set_xlim(0, 1.2)
ax.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
ax.set_ylim(-0.6, len(cuts) - 0.4)
ax.set_xlabel('value on the 20 patients')
ax.grid(axis='x', color=GRAY, alpha=0.2, lw=0.5)
for side in ['top', 'right', 'left']:
    ax.spines[side].set_visible(False)
ax.tick_params(axis='y', length=0)
ax.legend(fontsize=9.5, loc='upper center', bbox_to_anchor=(0.45, -0.16), ncol=2, frameon=False,
          handletextpad=0.3, columnspacing=1.2)
ax.set_title('The mean ranks "flag all 20" first; F1 ranks the threshold 0.5 first', fontsize=11)

plt.tight_layout()

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

plt.close()