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

# One picture for the whole lesson: every model family slides along the same complexity axis.
# Training error only falls; validation error is a U; the left side is bias, the right side variance.
# Under the axis, which way each family's knob must turn to move RIGHT (more complex).
# The curves are schematic (shape only), not data.
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'})
fig, (ax, kx) = plt.subplots(2, 1, figsize=(10, 6.6), gridspec_kw={'height_ratios': [2.3, 1]})

c = np.linspace(0, 1, 200)
train = 0.08 + 0.75 * np.exp(-3.2 * c)
val = 0.25 + 0.6 * np.exp(-4.5 * c) + 0.62 * c ** 3
best = c[np.argmin(val)]
ax.axvspan(0, best - 0.07, color=ORANGE, alpha=0.08, lw=0)
ax.axvspan(best + 0.07, 1, color=BLUE, alpha=0.08, lw=0)
ax.plot(c, train, color=GREEN, lw=2.5, label='training error: only goes down')
ax.plot(c, val, color=RED, lw=2.5, label='validation error: a U')
ax.axhline(0.22, color='gray', ls=':', lw=1.2)
ax.text(0.99, 0.2, 'noise floor (irreducible)', ha='right', va='top', fontsize=8.5)
ax.axvline(best, color='gray', ls='--', lw=1)
ax.annotate('the gap: a symptom of variance', xy=(0.885, 0.36), xytext=(0.885, 0.36), fontsize=9, ha='right', va='center')
ax.annotate('', xy=(0.9, float(np.interp(0.9, c, val))), xytext=(0.9, float(np.interp(0.9, c, train))),
            arrowprops=dict(arrowstyle='<->', color='gray', lw=1.2))
ax.text(0.04, 0.97, 'UNDERFITTING\nhigh bias', fontsize=10, color=ORANGE, fontweight='bold', va='top')
ax.text(0.96, 0.97, 'OVERFITTING\nhigh variance', fontsize=10, color=BLUE, fontweight='bold', va='top',
        ha='right')
ax.text(best, 0.97, 'sweet\nspot', fontsize=9.5, ha='center', va='top', fontweight='bold')
ax.set_xlim(0, 1); ax.set_ylim(0, 1.02)
ax.set_xticks([]); ax.set_yticks([])
ax.set_xlabel('model complexity  →', fontsize=10)
ax.set_ylabel('error', fontsize=10)
ax.legend(loc='upper center', bbox_to_anchor=(0.45, 0.8), frameon=False, fontsize=9)
for s in ['top', 'right']:
    ax.spines[s].set_visible(False)

knobs = [('polynomial', 'degree ↑'), ('Ridge / Lasso', 'alpha ↓'), ('logistic, SVM', 'C ↑'),
         ('k-NN', 'n_neighbors ↓'), ('tree', 'max_depth ↑'), ('tree', 'min_samples_leaf ↓')]
kx.axis('off')
kx.set_xlim(0, 6); kx.set_ylim(0, 1)
kx.text(3, 0.93, 'to move RIGHT (more complex), turn the knob:', ha='center', fontsize=10, fontweight='bold')
for i, (family, knob) in enumerate(knobs):
    reverse = '↓' in knob
    color = RED if reverse else GREEN
    kx.text(i + 0.5, 0.52, knob, ha='center', fontsize=9.5, color=color, fontweight='bold')
    kx.text(i + 0.5, 0.3, family, ha='center', fontsize=9)
kx.text(3, 0.03, 'green: bigger means more complex · red: the knobs that turn "backwards" '
        '(bigger means simpler)', ha='center', fontsize=8.5, style='italic')

plt.tight_layout()

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

plt.close()