import matplotlib.pyplot as plt
from io import StringIO

# The lesson's diagnosis as a flowchart: two questions about the training and validation errors
# decide whether a model suffers from bias, from variance, or is done, and each outcome lists
# what helps and what does not. Same content as the table at the end of the learning-curves section.
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(10, 6.6))
ax.set_xlim(0, 10); ax.set_ylim(-0.1, 6.4); ax.axis('off')


def box(x, y, text, color='none', w=None, size=9.5, bold=False):
    style = dict(boxstyle='round,pad=0.55', fc=color, ec='gray', lw=1.1,
                 alpha=0.22 if color != 'none' else 1)
    ax.text(x, y, text, ha='center', va='center', fontsize=size, color='gray', bbox=style,
            fontweight='bold' if bold else 'normal', linespacing=1.45)


def arrow(xa, ya, xb, yb, label='', side='left'):
    ax.annotate('', xy=(xb, yb), xytext=(xa, ya), arrowprops=dict(arrowstyle='-|>', color='gray', lw=1.3))
    if label:
        dx = -0.12 if side == 'left' else 0.12
        ax.text((xa + xb) / 2 + dx, (ya + yb) / 2, label, fontsize=9.5, fontweight='bold',
                ha='right' if side == 'left' else 'left', va='center')


box(5, 6.0, 'Cross-validate: read the training AND the validation error', bold=True)
box(3.2, 4.85, 'Is the TRAINING error already\ntoo high for the application?')
box(6.8, 3.75, 'Is VALIDATION much worse\nthan training?')
box(1.6, 2.35, 'HIGH BIAS\n(underfitting)\n\ntraining ≈ validation,\nboth too high', ORANGE, bold=True)
box(5.0, 2.35, 'HIGH VARIANCE\n(overfitting)\n\ntraining ≪ validation,\na large gap', BLUE, bold=True)
box(8.5, 2.35, 'GOOD ENOUGH\n\nstop tuning;\ntest once', GREEN, bold=True)

arrow(4.6, 5.72, 3.6, 5.28)
arrow(2.5, 4.45, 1.75, 3.25, 'yes', 'left')
arrow(4.25, 4.6, 6.0, 4.05, 'no', 'right')
arrow(6.2, 3.35, 5.3, 3.2, 'yes', 'left')
arrow(7.4, 3.35, 8.3, 3.05, 'no', 'right')

remedies = [(1.6, ORANGE, '✓ more capacity: higher degree,\n   larger C, smaller k, deeper tree\n'
             '✓ better features\n✓ less regularization\n✗ more data: the curves stay flat'),
            (5.0, BLUE, '✓ more data, if the gap still closes\n✓ more regularization: larger α,\n'
             '   smaller C, larger k, larger leaves\n✓ fewer features · average many models\n✗ more capacity'),
            (8.5, GREEN, '✓ prefer the simpler of\n   tied candidates\n✓ touch the test set\n   exactly once')]
for x, color, text in remedies:
    ax.text(x - 1.45, 1.25, text, fontsize=8.6, va='top', ha='left', color=color, linespacing=1.5)

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True, bbox_inches='tight', pad_inches=0.05)
print(buffer.getvalue())

plt.close()