import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

# The same experiment as the lesson's bias_variance(): 500 training sets of 30
# points from y = sin(2*pi*x) + noise (sd 0.3). Each panel draws 20 of the fitted
# curves, the average of all 500, and the truth.
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

f = lambda x: np.sin(2 * np.pi * x)
x_grid = np.linspace(0, 1, 101)
rng = np.random.default_rng(0)
sets = []
for _ in range(500):
    x = rng.uniform(0, 1, 30)
    sets.append((x, f(x) + rng.normal(0, 0.3, 30)))
    rng.normal(0, 0.3, x_grid.size)          # the fresh test noise drawn in bias_variance()

fig, axes = plt.subplots(1, 3, figsize=(10, 3.2), sharey=True)
for ax, d in zip(axes, [1, 3, 12]):
    preds = np.array([make_pipeline(PolynomialFeatures(d), StandardScaler(), LinearRegression())
                      .fit(x.reshape(-1, 1), y).predict(x_grid.reshape(-1, 1)) for x, y in sets])
    bias2 = np.mean((preds.mean(axis=0) - f(x_grid)) ** 2)
    var = np.mean(preds.var(axis=0))
    for p in preds[:20]:
        ax.plot(x_grid, p, color='#58a6ff', lw=0.8, alpha=0.45)
    ax.plot(x_grid, f(x_grid), '--', color='gray', lw=1.6, label='truth  sin(2πx)')
    ax.plot(x_grid, preds.mean(axis=0), color='#f0883e', lw=2.2, label='average of 500 fits')
    ax.set_ylim(-2, 2)
    ax.set_xlim(0, 1)
    ax.set_title(f'degree {d}:  bias² {bias2:.3f} · variance {var:.3f}', fontsize=9.5)
    ax.set_xlabel('x', fontsize=9)
    ax.tick_params(labelsize=8)
axes[0].plot([], [], color='#58a6ff', lw=0.8, label='20 of the 500 fits')
axes[0].legend(fontsize=7.5, loc='lower left', frameon=False)
axes[0].set_ylabel('prediction', fontsize=9)
axes[0].text(0.5, 1.72, 'too rigid: wrong in the same way every time', ha='center', fontsize=7.5, color='gray')
axes[1].text(0.5, 1.72, 'right shape, small wobble', ha='center', fontsize=7.5, color='gray')
axes[2].text(0.5, 1.72, 'right on average, wild at the edges', ha='center', fontsize=7.5, color='gray')

plt.tight_layout()

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

plt.close()