import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.ensemble import BaggingRegressor
from sklearn.tree import DecisionTreeRegressor

# The lesson's bias-variance experiment, drawn: 100 training sets of 100 points from
# y = sin(x) + 0.5 sin(3x) + noise (sd 0.3); each model is refit on every training set.


def truth(x):
    return np.sin(x) + 0.5 * np.sin(3 * x)


x_grid = np.linspace(0.2, 6.1, 100).reshape(-1, 1)     # where bias and variance are measured
x_fine = np.linspace(0, 2 * np.pi, 500).reshape(-1, 1) # where the curves are drawn
models = {'one tree, max_depth=3': DecisionTreeRegressor(max_depth=3),
          'one tree, full depth': DecisionTreeRegressor(),
          'bagging: 100 full trees': BaggingRegressor(DecisionTreeRegressor(), n_estimators=100,
                                                      random_state=0)}
preds = {name: [] for name in models}
curves = {name: [] for name in models}
for s in range(100):
    r = np.random.default_rng(s)
    x = r.uniform(0, 2 * np.pi, (100, 1))
    y = truth(x[:, 0]) + r.normal(0, 0.3, 100)
    for name, model in models.items():
        model.fit(x, y)
        preds[name].append(model.predict(x_grid))
        if s < 20:
            curves[name].append(model.predict(x_fine))

colors = ['#3fb950', '#f0883e', '#58a6ff']
fig, axes = plt.subplots(1, 3, figsize=(10, 3.3), sharey=True)
for ax, (name, P), col in zip(axes, preds.items(), colors):
    P = np.array(P)
    bias2 = ((P.mean(axis=0) - truth(x_grid[:, 0]))**2).mean()
    var = P.var(axis=0).mean()
    for c in curves[name]:
        ax.plot(x_fine[:, 0], c, color=col, lw=0.7, alpha=0.35)
    ax.plot(x_fine[:, 0], truth(x_fine[:, 0]), color='gray', lw=1.6, ls='--', label='true function')
    ax.set_title(f'{name}\nbias² {bias2:.4f} · variance {var:.4f}', fontsize=9)
    ax.set_xlabel('x', fontsize=8)
    ax.tick_params(labelsize=8)
    ax.set_ylim(-2.3, 2.3)
axes[0].set_ylabel('prediction', fontsize=8)
axes[0].legend(fontsize=7.5, loc='lower left')
fig.suptitle('20 of the 100 training sets: each faint line is the model refit on one of them', fontsize=9.5)

plt.tight_layout()

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

plt.close()