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

# F0.5, F1 and F2 on the 20 patients as the threshold moves (section 9).
# Bottom strip: the patients themselves; each step of the curves is the threshold crossing one of them.
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 f_beta(p, r, beta):
    return (1 + beta**2) * p * r / (beta**2 * p + r) if p + r > 0 else 0.0


def f_at(t, beta):
    pred = score >= t
    tp = (pred & (y_true == 1)).sum()
    return f_beta(tp / pred.sum(), tp / 10, beta) if tp else 0.0


fig, (ax, strip) = plt.subplots(2, 1, figsize=(8.4, 5.6), sharex=True,
                                gridspec_kw={'height_ratios': [4, 0.9], 'hspace': 0.08})
grid = np.linspace(0.0, 0.95, 1901)
curves = [(0.5, BLUE, 'F0.5: a false alarm weighs 4× a miss'), (1, GREEN, 'F1: both weigh the same'),
          (2, ORANGE, 'F2: a miss weighs 4× a false alarm')]
for beta, color, name in curves:
    best = max(score[::-1], key=lambda t: f_at(t, beta))      # flat between patients: the best is at a score
    ax.plot(grid, [f_at(t, beta) for t in grid], color=color, lw=2,
            label=f'{name}  ·  best t = {best:.2f} (dot)')
    ax.scatter([best], [f_at(best, beta)], color=color, s=70, zorder=5, edgecolor='white', lw=1)
ax.axvline(0.5, color=GRAY, ls=':', lw=1)
ax.text(0.507, 0.95, 'default 0.5', fontsize=8.5, color=GRAY)
ax.set_ylabel('F score on the 20 patients')
ax.set_ylim(0.0, 1.02)
ax.legend(fontsize=8.5, loc='lower left', frameon=False, labelcolor=GRAY)
ax.set_title('The larger β, the more recall counts, and the lower the best threshold', fontsize=10)
ax.grid(color=GRAY, alpha=0.2, lw=0.5)

# the patients on the same axis
for s, y in zip(score, y_true):
    strip.scatter([s], [0], s=55, color=GREEN if y else 'none', edgecolor=GREEN if y else RED, lw=1.5, zorder=3)
strip.scatter([], [], s=40, color=GREEN, label='sick patient: crossing it raises recall')
strip.scatter([], [], s=40, color='none', edgecolor=RED, lw=1.5, label='healthy patient: crossing it lowers precision')
strip.axvline(0.5, color=GRAY, ls=':', lw=1)
strip.set_yticks([])
strip.set_ylim(-1, 1)
fig.legend(*strip.get_legend_handles_labels(), fontsize=8.5, loc='lower center', ncol=2, frameon=False,
           labelcolor=GRAY, handletextpad=0.3)
strip.set_xlabel('threshold t   (flag every patient with score ≥ t)')
strip.set_xlim(0, 1)
for a in (ax, strip):
    for side in ['top', 'right']:
        a.spines[side].set_visible(False)
strip.spines['left'].set_visible(False)

fig.subplots_adjust(left=0.09, right=0.98, top=0.93, bottom=0.2, hspace=0.08)

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

plt.close()