import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.tree import DecisionTreeRegressor


# A noisy sine: 200 training points, 5,000 fresh points to measure the test error.
# The noise has variance 0.25, so no model can go below a test MSE of about 0.25.
def sample(n, seed):
    r = np.random.default_rng(seed)
    X = r.uniform(0, 6, (n, 1))
    return X, 2 * np.sin(X[:, 0]) + r.normal(0, 0.5, n)


X, y = sample(200, 0)
X_new, y_new = sample(5000, 1)
xs = np.linspace(0, 6, 600).reshape(-1, 1)

panels = [
    ('one stump', DecisionTreeRegressor(max_depth=1), '#f0883e'),
    ('forest: 200 stumps averaged', RandomForestRegressor(200, max_depth=1, random_state=0), '#f0883e'),
    ('boosting: 200 stumps added', GradientBoostingRegressor(n_estimators=200, max_depth=1,
                                                             learning_rate=0.1, random_state=0), '#3fb950'),
    ('forest: 200 deep trees averaged', RandomForestRegressor(200, random_state=0), '#58a6ff'),
]

fig, axes = plt.subplots(1, 4, figsize=(10, 2.9), sharey=True)
for ax, (title, model, color) in zip(axes, panels):
    model.fit(X, y)
    mse = np.mean((model.predict(X_new) - y_new) ** 2)
    ax.scatter(X, y, s=6, color='gray', alpha=0.45, lw=0)
    ax.plot(xs, 2 * np.sin(xs.ravel()), '--', color='gray', lw=1)
    ax.plot(xs, model.predict(xs), color=color, lw=2)
    ax.set_title(f'{title}\ntest MSE {mse:.2f}', fontsize=8.5)
    ax.set_xticks([])
    ax.tick_params(labelsize=7.5)
axes[0].set_ylabel('y', fontsize=8)
axes[1].text(3, -3.3, 'averaging does not\nremove bias', ha='center', fontsize=7.5, color='gray')
axes[2].text(3, -3.3, 'each stump fixes\nwhat is left', ha='center', fontsize=7.5, color='gray')
axes[3].text(3, -3.3, 'averaging removes\nvariance', ha='center', fontsize=7.5, color='gray')
axes[0].set_ylim(-3.9, 3.4)
fig.suptitle('Same 200 points, noise floor MSE 0.25 — the dashed line is the true curve', fontsize=9.5)

plt.tight_layout()

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

plt.close()