import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.impute import SimpleImputer
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

# The lesson's toy column: what fit, predict and the fit_transform bug do, step by step.
toy_train = pd.DataFrame({'total_bedrooms': [200, np.nan, 400, 600, 800]})
toy_new = pd.DataFrame({'total_bedrooms': [np.nan, 900]})
pipe = make_pipeline(SimpleImputer(strategy='median'), StandardScaler()).fit(toy_train)
imp, sc = pipe[0], pipe[1]
med, mu, sd = imp.statistics_[0], sc.mean_[0], sc.scale_[0]
tr_imp = imp.transform(toy_train).ravel()
tr_out = pipe.transform(toy_train).ravel()
new_imp = imp.transform(toy_new).ravel()
new_out = pipe.transform(toy_new).ravel()
bad = make_pipeline(SimpleImputer(strategy='median'), StandardScaler()).fit(toy_new)
b_med, b_mu = bad[0].statistics_[0], bad[1].mean_[0]
b_imp = bad[0].transform(toy_new).ravel()
b_out = bad.transform(toy_new).ravel()


def fmt(v):
    return ', '.join('NaN' if np.isnan(x) else f'{x:g}' for x in v)


BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(10, 4.2))
ax.set_xlim(0, 10); ax.set_ylim(0, 6.3)
ax.axis('off')

cols = [1.05, 3.65, 6.35, 8.95]
heads = ['data in', 'SimpleImputer(median)', 'StandardScaler', 'what the model gets']
for x, h in zip(cols, heads):
    ax.text(x, 6.05, h, ha='center', fontsize=9, fontweight='bold')


def box(x, y, text, color, w=2.3):
    ax.text(x, y, text, ha='center', va='center', fontsize=8.3, color='gray', linespacing=1.35,
            bbox=dict(boxstyle='round,pad=0.4', fc=color, ec=color, alpha=0.18))


def arrow(x0, x1, y):
    ax.annotate('', xy=(x1, y), xytext=(x0, y), arrowprops=dict(arrowstyle='->', color='gray', lw=1.1))


lanes = [
    (5.0, 'pipe.fit(train)', BLUE,
     [f'5 training rows\n{fmt(toy_train.total_bedrooms)}',
      f'LEARNS median = {med:g}\n→ {fmt(tr_imp)}',
      f'LEARNS mean = {mu:g}, std = {sd:g}\n→ {fmt(tr_out)}',
      'model.fit(…)']),
    (3.0, 'pipe.predict(new)', GREEN,
     [f'2 new rows\n{fmt(toy_new.total_bedrooms)}',
      f'reuses median {med:g}\n→ {fmt(new_imp)}',
      f'reuses (x − {mu:g}) / {sd:g}\n→ {fmt(new_out)}',
      f'{fmt(new_out)}\n900 is 2 std above the\ntraining average']),
    (0.9, 'bug 1: fit_transform(new)', RED,
     [f'the same 2 rows\n{fmt(toy_new.total_bedrooms)}',
      f're-learns median = {b_med:g}\n→ {fmt(b_imp)}',
      f're-learns mean = {b_mu:g}, std = 0\n→ {fmt(b_out)}',
      f'{fmt(b_out)}\nboth rows now look\nperfectly average']),
]
for y, label, color, texts in lanes:
    ax.text(0.0, y + 0.72, label, fontsize=8.8, color=color, fontweight='bold', family='monospace')
    for x, t in zip(cols, texts):
        box(x, y, t, color)
    for a, b in zip(cols[:-1], cols[1:]):
        arrow(a + 1.2, b - 1.2, y)
# the learned state flows from fit to predict
for x in cols[1:3]:
    ax.annotate('', xy=(x, 3.55), xytext=(x, 4.45),
                arrowprops=dict(arrowstyle='->', color=BLUE, lw=1.2, ls='--'))
ax.text(5.0, 4.0, 'learned state, frozen', ha='center', va='center', fontsize=8, color=BLUE)

plt.tight_layout()

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

plt.close()