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

# Resampling before vs after the split (section 10). The numbers are the 5-NN with
# RandomOverSampler from the lesson's code: recall / precision on each "test" set.
RED, GREEN, BLUE, GRAY = '#ff7b72', '#3fb950', '#58a6ff', 'gray'

fig, ax = plt.subplots(figsize=(9.6, 5.6))
ax.set_xlim(0, 12)
ax.set_ylim(0, 7.4)
ax.axis('off')


def box(x, y, w, h, text, color, size=9):
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle='round,pad=0.02,rounding_size=0.12',
                                fc=color, ec='none', alpha=0.13))
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle='round,pad=0.02,rounding_size=0.12',
                                fc='none', ec=color, lw=1.3))
    ax.text(x + w / 2, y + h / 2, text, ha='center', va='center', fontsize=size, color=GRAY)


def arrow(x0, y0, x1, y1, color=GRAY, text=None, dx=0, dy=0.12):
    ax.add_patch(FancyArrowPatch((x0, y0), (x1, y1), arrowstyle='-|>', mutation_scale=13, color=color, lw=1.2))
    if text:
        ax.text((x0 + x1) / 2 + dx, (y0 + y1) / 2 + dy, text, ha='center', va='bottom', fontsize=7.8,
                color=GRAY, style='italic')


# wrong: resample everything, then split
ax.text(0.2, 7.05, 'Wrong: resample, then split', fontsize=11, fontweight='bold', color=RED, va='top')
box(0.2, 4.9, 2.1, 1.2, 'all rows\n2% fraud', BLUE)
box(2.9, 4.9, 2.5, 1.2, 'oversample\nevery fraud ×50', RED)
box(6.0, 4.9, 1.4, 1.2, 'split', GRAY)
box(8.0, 4.9, 2.3, 1.2, '"test": 50% fraud,\ncopies of the\ntraining frauds', RED, size=8.5)
arrow(2.32, 5.5, 2.88, 5.5)
arrow(5.42, 5.5, 5.98, 5.5)
arrow(7.42, 5.5, 7.98, 5.5)
ax.text(10.55, 5.82, '5-NN', fontsize=9, color=GRAY, va='center')
ax.text(10.55, 5.35, 'recall 1.00\nprecision 0.96', fontsize=9, color=RED, va='center', fontweight='bold')
ax.text(10.55, 4.8, 'too good to be true', fontsize=8, color=GRAY, va='center', style='italic')

ax.plot([0.2, 11.8], [4.15, 4.15], color=GRAY, lw=0.6, alpha=0.4)

# right: split first, resample only inside the pipeline's fit
ax.text(0.2, 3.85, 'Right: split, then resample inside the pipeline', fontsize=11, fontweight='bold',
        color=GREEN, va='top')
box(0.2, 1.35, 2.1, 1.2, 'all rows\n2% fraud', BLUE)
box(2.9, 1.35, 1.4, 1.2, 'split', GRAY)
box(4.9, 2.2, 4.1, 0.95, 'training rows (2% fraud)\n→ oversample → fit', GREEN, size=8.5)
box(4.9, 0.3, 4.1, 0.95, 'validation / test (2% fraud, untouched)\n→ predict: the sampler is skipped', GREEN,
    size=8.5)
arrow(2.32, 1.95, 2.88, 1.95)
arrow(4.32, 2.2, 4.88, 2.6)
arrow(4.32, 1.7, 4.88, 0.85)
arrow(6.95, 2.18, 6.95, 1.27, GREEN, 'fitted model', dx=0.65, dy=-0.12)
ax.text(9.4, 1.13, '5-NN', fontsize=9, color=GRAY, va='center')
ax.text(9.4, 0.66, 'recall 0.42\nprecision 0.20', fontsize=9, color=GREEN, va='center', fontweight='bold')
ax.text(9.4, 0.11, 'what production will see', fontsize=8, color=GRAY, va='center', style='italic')

plt.tight_layout()

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

plt.close()