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

# The 20 referred patients of the metrics lesson, sorted by the model's 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])
P, N = sick.sum(), (1 - sick).sum()

# Lower the threshold one patient at a time: a sick patient moves the point up by 1/P,
# a healthy one moves it right by 1/N
fpr, tpr = [0.0], [0.0]
for s in sick:
    fpr.append(fpr[-1] + (1 - s) / N)
    tpr.append(tpr[-1] + s / P)
fpr, tpr = np.array(fpr), np.array(tpr)
auc = roc_auc_score(sick, score)

fig, ax = plt.subplots(figsize=(6.5, 5.2))
ax.fill_between(fpr, tpr, step='post', color='#58a6ff', alpha=0.12, lw=0)
ax.plot([0, 1], [0, 1], '--', color='gray', lw=1)
ax.text(0.76, 0.69, 'random scorer (AUC = 0.50)', color='gray', fontsize=8, rotation=38.5,
        ha='center', va='center')
for i, s in enumerate(sick):
    color = '#3fb950' if s else '#ff7b72'
    ax.plot(fpr[i:i + 2], tpr[i:i + 2], color=color, lw=2.4, solid_capstyle='round')
    xm, ym = (fpr[i] + fpr[i + 1]) / 2, (tpr[i] + tpr[i + 1]) / 2
    if s:
        ax.text(xm - 0.012, ym, f'P{i + 1:02d}', fontsize=7.5, color='#3fb950', ha='right', va='center')
    else:
        ax.text(xm, ym + 0.012, f'P{i + 1:02d}', fontsize=7.5, color='#ff7b72', ha='center', va='bottom')
ax.plot(fpr, tpr, 'o', color='gray', ms=2.5)

# the operating point of the metrics lesson: threshold 0.5 (after P09)
ax.plot(0.2, 0.7, 'o', ms=10, mfc='none', mec='#f0883e', mew=2)
ax.annotate('threshold 0.5\nTP = 7, FP = 2\nTPR = 0.7, FPR = 0.2', xy=(0.2, 0.7), xytext=(-0.05, 0.9),
            fontsize=8.5, color='#f0883e', arrowprops=dict(arrowstyle='->', color='#f0883e'))
ax.annotate('(0, 0): threshold above\nevery score, nobody flagged', xy=(0, 0), xytext=(0.12, 0.06),
            fontsize=8, color='gray', arrowprops=dict(arrowstyle='->', color='gray'))
ax.annotate('(1, 1): everybody flagged', xy=(1, 1), xytext=(0.52, 0.86),
            fontsize=8, color='gray', arrowprops=dict(arrowstyle='->', color='gray'))
ax.text(0.55, 0.16, f'shaded area = AUC = {auc:.2f}', fontsize=9, color='#58a6ff')
ax.plot([], [], color='#3fb950', lw=2.4, label='sick patient: one step up (1/10)')
ax.plot([], [], color='#ff7b72', lw=2.4, label='healthy patient: one step right (1/10)')
ax.legend(fontsize=8, loc='lower right', bbox_to_anchor=(1.0, 0.2))
ax.set_xlim(-0.07, 1.02)
ax.set_ylim(-0.02, 1.04)
ax.set_aspect('equal')
ax.set_xticks(np.arange(0, 1.01, 0.1))
ax.set_yticks(np.arange(0, 1.01, 0.1))
ax.grid(color='gray', alpha=0.2, lw=0.5)
ax.set_xlabel('false positive rate (FPR) = FP / 10 healthy')
ax.set_ylabel('true positive rate (TPR, recall) = TP / 10 sick')
ax.set_title('ROC of the 20 patients, built by lowering the threshold one patient at a time', fontsize=9.5)

plt.tight_layout()

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

plt.close()