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

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 lesson's running example: 20 patients, the model's score, and the truth
ids = [f"P{i:02d}" for i in range(1, 21)]
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])
y_pred = (score >= 0.5).astype(int)
color = {(1, 1): GREEN, (0, 1): ORANGE, (1, 0): RED, (0, 0): BLUE}
name = {(1, 1): 'TP', (0, 1): 'FP', (1, 0): 'FN', (0, 0): 'TN'}

fig, (ax, axm) = plt.subplots(1, 2, figsize=(10, 3.6), gridspec_kw={'width_ratios': [2.6, 1]})

# --- left: the patients on the score axis ------------------------------------
ax.axvspan(0.5, 1.02, color=GRAY, alpha=0.10, lw=0)
ax.axvline(0.5, color=GRAY, ls='--', lw=1.2)
ax.text(0.76, 1.72, 'score ≥ 0.5 → model says "sick"', ha='center', fontsize=8.5)
ax.text(0.25, 1.72, 'score < 0.5 → model says "healthy"', ha='center', fontsize=8.5)
for s, t, p, i in zip(score, y_true, y_pred, ids):
    row = 1.0 if t == 1 else 0.0
    ax.scatter(s, row, s=90, color=color[t, p], zorder=3, edgecolors='none')
    ax.text(s, row + 0.17 if t == 1 else row - 0.17, i, rotation=90, fontsize=7.5,
            ha='center', va='bottom' if t == 1 else 'top')
ax.set_yticks([0, 1], ['healthy\n(10)', 'sick\n(10)'])
ax.set_ylim(-0.75, 1.95)
ax.set_xlim(-0.02, 1.02)
ax.set_xlabel("model's predicted probability of disease (score)")
for side in ['top', 'right', 'left']:
    ax.spines[side].set_visible(False)
ax.tick_params(axis='y', length=0)
for key, (x, y) in {(1, 1): (0.6, 1.42), (1, 0): (0.02, 1.42),
                    (0, 1): (0.6, -0.62), (0, 0): (0.02, -0.62)}.items():
    ax.text(x, y, {(1, 1): 'TP: sick, flagged', (1, 0): 'FN: sick, missed',
                   (0, 1): 'FP: healthy, flagged', (0, 0): 'TN: healthy, cleared'}[key],
            color=color[key], fontsize=8.5, weight='bold')

# --- right: the confusion matrix in the standard layout ----------------------
# rows = actual condition, columns = predicted condition, positive (sick) first:
# TP top-left, FN top-right, FP bottom-left, TN bottom-right. scikit-learn's default puts
# the labels in ascending order (0 first, TN top-left); labels=[1, 0] gives this layout.
axm.set_xlim(0, 2); axm.set_ylim(0, 2); axm.invert_yaxis(); axm.set_aspect('equal')
for t in (0, 1):
    for p in (0, 1):
        n = int(((y_true == t) & (y_pred == p)).sum())
        r, c = 1 - t, 1 - p
        axm.add_patch(Rectangle((c + 0.04, r + 0.04), 0.92, 0.92, color=color[t, p], alpha=0.85, lw=0))
        axm.text(c + 0.5, r + 0.38, name[t, p], ha='center', va='center', fontsize=10, color='white')
        axm.text(c + 0.5, r + 0.66, str(n), ha='center', va='center', fontsize=16, color='white',
                 weight='bold')
axm.set_xticks([0.5, 1.5], ['sick', 'healthy'])
axm.set_yticks([0.5, 1.5], ['sick', 'healthy'])
axm.set_xlabel('predicted condition')
axm.set_ylabel('actual condition')
axm.xaxis.set_label_position('top'); axm.xaxis.tick_top()
axm.tick_params(length=0, labelsize=8.5)
for s in axm.spines.values():
    s.set_visible(False)
axm.set_title('confusion_matrix(y_true, y_pred,\n                 labels=[1, 0])', fontsize=7.5, y=-0.2, family='monospace')

fig.suptitle('Threshold 0.5 sorts the 20 patients into four cells', fontsize=11)
plt.tight_layout()

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

plt.close()