import matplotlib.pyplot as plt
from io import StringIO
from matplotlib.patches import FancyBboxPatch

# The model-selection workflow of the lesson's lab, as five stages. Stages 2 to 4 loop on the
# training data only; the test set sits behind a wall that is crossed once, at the end.
BLUE, ORANGE, GREEN, RED, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72', '#bc8cff'
plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(11, 4.3))
ax.set_xlim(0, 11); ax.set_ylim(0, 4.3); ax.axis('off')

stages = [
    (0.15, RED, '1 · Lock the test', 'split once;\nnever look at it\nwhile choosing'),
    (2.35, PURPLE, '2 · One cv for all', 'the same folds\nfor every model,\nso scores compare'),
    (4.55, ORANGE, '3 · Diagnose', 'validation and\nlearning curves:\nbias or variance?'),
    (6.75, BLUE, '4 · Treat & select', 'turn the knob, try\nfamilies; on a tie,\nkeep the simpler'),
    (8.95, GREEN, '5 · Test once', 'report this\nnumber; do not\ngo back'),
]
W, H, Y = 1.9, 2.0, 1.35
for x, color, title, text in stages:
    ax.add_patch(FancyBboxPatch((x, Y), W, H, boxstyle='round,pad=0.04,rounding_size=0.12',
                                fc=color, ec=color, alpha=0.18, lw=1.5))
    ax.add_patch(FancyBboxPatch((x, Y), W, H, boxstyle='round,pad=0.04,rounding_size=0.12',
                                fc='none', ec=color, lw=1.5))
    ax.text(x + W / 2, Y + H - 0.3, title, ha='center', va='center', fontsize=9.5, fontweight='bold',
            color=color)
    ax.text(x + W / 2, Y + H / 2 - 0.2, text, ha='center', va='center', fontsize=9, linespacing=1.45)
for x, _, _, _ in stages[:-1]:
    ax.annotate('', xy=(x + W + 0.27, Y + H / 2), xytext=(x + W + 0.05, Y + H / 2),
                arrowprops=dict(arrowstyle='-|>', color='gray', lw=1.4))

# loop: diagnose <-> treat, on training data only
ax.annotate('', xy=(4.55 + W / 2, Y - 0.05), xytext=(6.75 + W / 2, Y - 0.05),
            arrowprops=dict(arrowstyle='-|>', color='gray', lw=1.3, connectionstyle='arc3,rad=-0.35'))
ax.text(6.65, 0.32, 'repeat, on the training data only', ha='center', fontsize=9, style='italic')
# the wall before the test set
ax.plot([8.83, 8.83], [0.95, 3.75], color=RED, lw=2.5, ls=(0, (4, 2)))
ax.text(8.83, 3.95, 'one-way door', ha='center', fontsize=9, color=RED, fontweight='bold')
ax.text(1.1, 3.95, 'everything left of the door uses only the training rows', ha='left', fontsize=9)

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

plt.close()