import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from scipy.stats import norm
from sklearn.metrics import roc_auc_score, roc_curve

# Healthy scores ~ N(0, 1), sick scores ~ N(d, 1): the further apart, the larger the AUC.
# For this model AUC = P(sick > healthy) = Phi(d / sqrt(2)) (same code as in the lesson).
rng = np.random.default_rng(0)
n = 2_000
seps = [0.5, 1.5, 3.0]
titles = ['heavy overlap', 'moderate overlap', 'well separated']
GREEN, RED, BLUE = '#3fb950', '#ff7b72', '#58a6ff'

fig, axes = plt.subplots(2, 3, figsize=(9.6, 6.2), gridspec_kw={'height_ratios': [1, 1.35]})
bins = np.linspace(-4, 7, 45)
for j, (d, title) in enumerate(zip(seps, titles)):
    neg, pos = rng.normal(0, 1, n), rng.normal(d, 1, n)
    y, s = np.r_[np.zeros(n), np.ones(n)], np.r_[neg, pos]
    auc = roc_auc_score(y, s)
    top, bottom = axes[0, j], axes[1, j]
    top.hist(neg, bins=bins, color=RED, alpha=0.45, label='healthy')
    top.hist(pos, bins=bins, color=GREEN, alpha=0.45, label='sick')
    top.set_title(f'{title}: d = {d}', fontsize=10)
    top.set_yticks([])
    top.set_xlabel('score', fontsize=9)
    if j == 0:
        top.legend(fontsize=8, loc='upper right', frameon=False)

    fpr, tpr, _ = roc_curve(y, s)
    bottom.plot(fpr, tpr, color=BLUE, lw=1.8)
    bottom.fill_between(fpr, tpr, color=BLUE, alpha=0.12, lw=0)
    bottom.plot([0, 1], [0, 1], '--', color='gray', lw=1)
    bottom.text(0.97, 0.06, f'AUC = {auc:.3f}\n$\\Phi(d/\\sqrt{{2}})$ = {norm.cdf(d / np.sqrt(2)):.3f}',
                ha='right', va='bottom', fontsize=9, color='gray')
    bottom.set_xlim(0, 1)
    bottom.set_ylim(0, 1.02)
    bottom.set_aspect('equal')
    bottom.set_xlabel('false positive rate', fontsize=9)
    if j == 0:
        bottom.set_ylabel('true positive rate', fontsize=9)
    bottom.grid(color='gray', alpha=0.2, lw=0.5)

for ax in axes.flat:
    ax.tick_params(colors='gray', labelsize=8)
    for spine in ax.spines.values():
        spine.set_color('gray')
    ax.xaxis.label.set_color('gray')
    ax.yaxis.label.set_color('gray')
    ax.title.set_color('gray')

plt.tight_layout()

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

plt.close()