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

BLUE, ORANGE, GREEN, GRAY = '#58a6ff', '#f0883e', '#3fb950', '#8b949e'

fig, ax = plt.subplots(figsize=(9, 4.8))
ax.set_xlim(0, 100)
ax.set_ylim(0, 50)
ax.axis('off')

# --- the dataset, cut into three blocks of rows ------------------------------
segments = [(0, 60, BLUE, 'training', '60%'),
            (60, 80, ORANGE, 'validation', '20%'),
            (80, 100, GREEN, 'test', '20%')]
for x0, x1, color, name, share in segments:
    ax.add_patch(FancyBboxPatch((x0 + 0.4, 38), x1 - x0 - 0.8, 7, boxstyle='round,pad=0,rounding_size=0.8',
                                facecolor=color, edgecolor='none', alpha=0.9))
    ax.text((x0 + x1) / 2, 41.5, f'{name}  ·  {share}', ha='center', va='center',
            fontsize=10, color='white', weight='bold')
# thin row lines: the blocks are rows of the same table
for x in range(2, 100, 2):
    ax.plot([x, x], [38.3, 44.7], color='white', lw=0.4, alpha=0.35)
ax.text(0, 47.5, 'all labeled data (rows shuffled; stratified by class)', fontsize=9, color=GRAY)

# --- what each block is for --------------------------------------------------
rows = [
    ('question it\nanswers', ['"what are the\nparameters?"', '"which model or\nsetting is best?"', '"how good on\nnew data?"']),
    ('used by', ['.fit() of every\ncandidate', 'the score of\nevery candidate', 'the final\nmodel only']),
    ('touched', ['at every fit', 'many times:\nit wears out', 'ONCE,\nat the end']),
]
centers = [30, 70, 90]
for i, (label, cells) in enumerate(rows):
    y = 30.5 - i * 7
    ax.text(-1, y, label, ha='right', va='center', fontsize=8, color=GRAY, style='italic')
    for cx, txt, (_, _, color, _, _) in zip(centers, cells, segments):
        ax.text(cx, y, txt, ha='center', va='center', fontsize=8.5, color=color if i == 2 else 'black',
                weight='bold' if i == 2 else 'normal', linespacing=1.1)
ax.plot([0, 100], [35, 35], color=GRAY, lw=0.5)
ax.plot([0, 100], [14.2, 14.2], color=GRAY, lw=0.5)

# --- the workflow ------------------------------------------------------------
def box(x, y, w, text, color):
    ax.add_patch(FancyBboxPatch((x, y), w, 5.2, boxstyle='round,pad=0.3,rounding_size=1',
                                facecolor='none', edgecolor=color, lw=1.4))
    ax.text(x + w / 2, y + 2.6, text, ha='center', va='center', fontsize=8.5, linespacing=1.1)

def arrow(p, q, color, rad=0.0, text=None, tx=0, ty=0):
    ax.add_patch(FancyArrowPatch(p, q, arrowstyle='-|>', mutation_scale=11, color=color, lw=1.3,
                                 connectionstyle=f'arc3,rad={rad}'))
    if text:
        ax.text(tx, ty, text, ha='center', va='center', fontsize=7.5, color=color)

Y = 1.8
box(3, Y, 20, 'fit a candidate\non training', BLUE)
box(33, Y, 20, 'score it on\nvalidation', ORANGE)
box(60, Y, 16, 'refit winner on\ntrain + validation', GRAY)
box(83, Y, 15, 'score once\non test', GREEN)
arrow((23.4, Y + 2.6), (32.6, Y + 2.6), GRAY)
arrow((33, Y + 5.6), (23, Y + 5.6), ORANGE, rad=0.55, text='try another model / hyperparameter', tx=28, ty=Y + 10.6)
arrow((53.4, Y + 2.6), (59.6, Y + 2.6), GRAY, text='best', tx=56.5, ty=Y + 4.3)
arrow((76.4, Y + 2.6), (82.6, Y + 2.6), GRAY)
ax.text(90.5, Y - 2.2, 'the number you report', ha='center', fontsize=7.5, color=GREEN)
ax.text(28, Y - 2.2, 'development loop: CV can replace this pair (next section)', ha='center',
        fontsize=7.5, color=GRAY)

plt.tight_layout(pad=0.3)

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

plt.close()