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

# The classic bias/variance dartboard. Each dot is the prediction of one model trained on a
# different training set; the bullseye is the truth. Bias = how far the average shot (the cross)
# lands from the bullseye; variance = how widely the shots scatter around their own average.
# The labels are the measured numbers of the lesson's polynomial experiment (500 training sets).
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

panels = [  # (row, col, center offset, spread, color, title, subtitle)
    (0, 0, (0.0, 0.0), 0.16, GREEN, 'low bias · low variance', 'degree 3: bias² 0.005, variance 0.018'),
    (0, 1, (0.0, 0.0), 0.62, BLUE, 'low bias · high variance', 'degree 12: bias² 0.001, variance 38.9'),
    (1, 0, (0.95, 0.55), 0.16, ORANGE, 'high bias · low variance', 'degree 1: bias² 0.204, variance 0.022'),
    (1, 1, (0.95, 0.55), 0.62, RED, 'high bias · high variance', 'the worst of both'),
]
rng = np.random.default_rng(3)
fig, axes = plt.subplots(2, 2, figsize=(8.2, 8.2))
for row, col, (cx, cy), spread, color, title, sub in panels:
    ax = axes[row, col]
    for r, a in zip([2.0, 1.5, 1.0, 0.5], [0.05, 0.08, 0.11, 0.16]):
        ax.add_patch(Circle((0, 0), r, fc='gray', alpha=a, ec='gray', lw=0.8))
    ax.add_patch(Circle((0, 0), 0.12, fc='gray', ec='none', alpha=0.6))
    shots = rng.normal([cx, cy], spread, size=(14, 2))
    shots = shots[np.hypot(*shots.T) < 1.95]
    ax.scatter(*shots.T, s=34, color=color, edgecolor='white', linewidth=0.6, zorder=3)
    mx, my = shots.mean(axis=0)
    ax.scatter([mx], [my], marker='X', s=150, color=color, edgecolor='gray', linewidth=1.0, zorder=4)
    if cx or cy:
        ax.annotate('', xy=(mx, my), xytext=(0, 0),
                    arrowprops=dict(arrowstyle='->', color='gray', lw=1.4, ls='--'))
        ax.text(mx / 2 - 0.15, my / 2 + 0.2, 'bias', fontsize=9, ha='right')
    ax.set_xlim(-2.1, 2.1); ax.set_ylim(-2.1, 2.1)
    ax.set_aspect('equal'); ax.axis('off')
    ax.set_title(title, fontsize=11, color=color, fontweight='bold', pad=4)
    ax.text(0, -2.3, sub, fontsize=9, ha='center', va='top')

fig.text(0.5, 0.995, 'Each dot: one model, trained on a different sample · ✕ their average · center: the truth',
         ha='center', va='top', fontsize=10)
fig.text(0.29, 0.955, 'shots close together', ha='center', fontsize=9, style='italic')
fig.text(0.75, 0.955, 'shots spread out', ha='center', fontsize=9, style='italic')
plt.tight_layout(rect=(0, 0, 1, 0.95), h_pad=3)

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

plt.close()