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})

# Triage of 200 patients into three risk levels (rows = truth, columns = prediction)
labels = ['low', 'medium', 'high']
counts = np.array([[110, 9, 1],
                   [8, 46, 6],
                   [2, 8, 10]])
support = counts.sum(axis=1)
precision = np.diag(counts) / counts.sum(axis=0)
recall = np.diag(counts) / support
f1 = 2 * precision * recall / (precision + recall)
macro, weighted = f1.mean(), (f1 * support).sum() / support.sum()

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4.0), gridspec_kw={'width_ratios': [1, 1.25]})

# --- left: the 3x3 matrix, with the "high" class read one-vs-rest -------------
ax1.set_xlim(0, 3); ax1.set_ylim(0, 3); ax1.invert_yaxis(); ax1.set_aspect('equal')
for i in range(3):
    for j in range(3):
        if i == j:
            face, alpha = GREEN, 0.85
        elif i == 2 or j == 2:
            face, alpha = RED if i == 2 else ORANGE, 0.55    # errors that involve "high"
        else:
            face, alpha = GRAY, 0.15
        ax1.add_patch(Rectangle((j + 0.04, i + 0.04), 0.92, 0.92, color=face, alpha=alpha, lw=0))
        ax1.text(j + 0.5, i + 0.5, str(counts[i, j]), ha='center', va='center', fontsize=13,
                 color='white' if (i == j or i == 2 or j == 2) else GRAY, weight='bold')
ax1.set_xticks([0.5, 1.5, 2.5], labels)
ax1.set_yticks([0.5, 1.5, 2.5], [f'{l} ({n})' for l, n in zip(labels, support)])
ax1.xaxis.tick_top(); ax1.xaxis.set_label_position('top')
ax1.set_xlabel('predicted risk')
ax1.set_ylabel('true risk (support)')
ax1.tick_params(length=0)
for s in ax1.spines.values():
    s.set_visible(False)
ax1.text(1.5, 3.25, 'for class "high": TP = 10 (green), FN = 2 + 8 (red row),\n'
         'FP = 1 + 6 (orange column), TN = everything else = 173',
         ha='center', va='top', fontsize=8)

# --- right: per-class F1 and the two averages ---------------------------------
x = np.arange(3)
ax2.bar(x, f1, width=support / support.max() * 0.8, color=[BLUE, BLUE, RED], alpha=0.85)
for xi, v, n in zip(x, f1, support):
    ax2.text(xi, v + 0.02, f'{v:.3f}', ha='center', fontsize=9, weight='bold')
ax2.axhline(weighted, color=GREEN, lw=1.6, ls='--')
ax2.axhline(macro, color=ORANGE, lw=1.6, ls='--')
ax2.text(2.45, weighted + 0.015, f'weighted avg {weighted:.3f}', color=GREEN, fontsize=8.5, va='bottom')
ax2.text(2.45, macro - 0.015, f'macro avg {macro:.3f}', color=ORANGE, fontsize=8.5, va='top')
ax2.set_xticks(x, [f'{l}\n{n} patients' for l, n in zip(labels, support)])
ax2.set_xlim(-0.6, 3.4)
ax2.set_ylim(0, 1.05)
ax2.set_ylabel('F1 of the class')
ax2.set_title('bar width = number of patients in the class', fontsize=8.5)
for side in ['top', 'right']:
    ax2.spines[side].set_visible(False)

fig.suptitle('The rare "high" class fails; the weighted average barely notices', fontsize=10.5)
plt.tight_layout()

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

plt.close()