import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.preprocessing import StandardScaler

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

train = np.array([[20.], [25.], [30.], [35.], [40.]])     # ages seen during training
test = np.array([[60.], [70.]])                           # two new, older passengers

right = StandardScaler().fit(train)
wrong = StandardScaler().fit(test)                        # the bug: fit again on the test rows
panels = [
    ('raw ages (years)', train.ravel(), test.ravel(), None),
    (f'RIGHT: test transformed with the TRAINING mean {right.mean_[0]:.0f} and std {right.scale_[0]:.2f}',
     right.transform(train).ravel(), right.transform(test).ravel(), '#3fb950'),
    (f'WRONG: scaler refit on the test rows (mean {wrong.mean_[0]:.0f}, std {wrong.scale_[0]:.0f})',
     right.transform(train).ravel(), wrong.transform(test).ravel(), '#ff7b72'),
]

fig, axes = plt.subplots(3, 1, figsize=(8.5, 4.6))
for ax, (title, tr, te, col) in zip(axes, panels):
    ax.axhline(0, color='gray', lw=0.6, alpha=0.4)
    ax.scatter(tr, np.zeros_like(tr), s=55, color='#58a6ff', zorder=3, edgecolor='white', lw=0.6,
               label='training rows')
    ax.scatter(te, np.zeros_like(te), s=70, color='#f0883e', zorder=4, edgecolor='white', lw=0.6,
               marker='D', label='test rows')
    for a, v in zip(train.ravel(), tr):
        ax.text(v, -0.55, f'{a:.0f}', fontsize=7.5, ha='center', color='#58a6ff')
    for a, v in zip(test.ravel(), te):
        ax.text(v, 0.4, f'{a:.0f} y → {v:.2f}' if col else f'{a:.0f}', fontsize=8, ha='center', color='#f0883e')
    ax.set_title(title, fontsize=9, loc='left', color=col or 'gray')
    ax.set_yticks([])
    ax.set_ylim(-0.9, 0.9)
    ax.tick_params(labelsize=8)
    for side in ['left', 'right', 'top']:
        ax.spines[side].set_visible(False)
axes[0].set_xlim(15, 75)
axes[0].legend(fontsize=8, loc='upper right', frameon=False, ncol=2, bbox_to_anchor=(1, 1.45))
for ax in axes[1:]:
    ax.set_xlim(-2, 6.2)
axes[2].text(4.2, -0.5, 'the 60-year-old now looks like a 23-year-old', fontsize=8.5, color='#ff7b72',
             ha='center')
plt.tight_layout()

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

plt.close()