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

# Where each remedy for imbalance acts on the way from training rows to a decision.
# ✓ = measured in this lesson.
BLUE, ORANGE, GREEN, GRAY = '#58a6ff', '#f0883e', '#3fb950', 'gray'

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


def box(x, y, w, h, color, alpha=0.12, lw=1.4):
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle='round,pad=0.02,rounding_size=0.15',
                                fc=color, ec=color, alpha=alpha, lw=0))
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle='round,pad=0.02,rounding_size=0.15',
                                fc='none', ec=color, lw=lw))


def arrow(x0, y0, x1, y1, color=GRAY):
    ax.add_patch(FancyArrowPatch((x0, y0), (x1, y1), arrowstyle='-|>', mutation_scale=14,
                                 color=color, lw=1.3))


# the pipeline, left to right
stages = [(0.3, 'training rows\n(2% positive)'), (3.3, 'fit: minimize\nthe loss'),
          (6.3, 'scores\n(probabilities)'), (9.3, 'threshold\n→ decision')]
for x, label in stages:
    box(x, 5.7, 2.4, 1.0, GRAY, alpha=0.06, lw=1.0)
    ax.text(x + 1.2, 6.2, label, ha='center', va='center', fontsize=9.5, color=GRAY)
for x in (2.7, 5.7, 8.7):
    arrow(x + 0.02, 6.2, x + 0.58, 6.2)

cards = [
    (0.1, BLUE, 'Data level', 'change the training rows', [
        'undersample:', '  RandomUnderSampler ✓', '  TomekLinks', '  EditedNearestNeighbours',
        'oversample:', '  RandomOverSampler ✓', '  SMOTE ✓', '  BorderlineSMOTE, ADASYN',
        'combine:', '  SMOTEENN, SMOTETomek'],
     'changes the model and\nits probabilities', 1.5),
    (4.1, ORANGE, 'Algorithm level', 'change the loss', [
        "  class_weight='balanced' ✓", '  sample_weight (a cost per row)',
        'balanced ensembles:', '  BalancedRandomForestClassifier', '  EasyEnsembleClassifier',
        '  RUSBoostClassifier'],
     'changes the model and\nits probabilities', 4.5),
    (8.1, GREEN, 'Decision level', 'change the threshold', [
        'chosen on validation by', 'best F1 ✓, a recall target ✓', 'or minimum cost ✓',
        'automated:', '  TunedThresholdClassifierCV', '',
        'no retraining: it can change', 'after deployment'],
     'keeps the model and\nits probabilities', 10.5),
]
for x, color, title, subtitle, lines, effect, target in cards:
    box(x, 0.25, 3.8, 4.75, color)
    ax.text(x + 0.2, 4.7, title, fontsize=11, fontweight='bold', color=color, va='top')
    ax.text(x + 0.2, 4.28, subtitle, fontsize=9, color=GRAY, va='top', style='italic')
    for i, line in enumerate(lines):
        mono = line.startswith('  ')
        ax.text(x + (0.4 if mono else 0.2), 3.8 - 0.27 * i, line.strip(), fontsize=8.4, color=GRAY,
                va='top', family='monospace' if mono else None)
    ax.text(x + 0.2, 0.42, effect, fontsize=8.6, color=color, va='bottom')
    arrow(x + 1.9, 5.05, target, 5.65, color)

ax.text(11.9, 0.02, '✓ measured in this lesson', ha='right', va='bottom', fontsize=8, color=GRAY)

plt.tight_layout()

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

plt.close()