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

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

# The lesson's bias_variance(): 500 training sets of 30 points from sin(2*pi*x) + noise (sd 0.3),
# the same sets for every model; bias^2 and variance averaged over 101 points of [0, 1].
f = lambda x: np.sin(2 * np.pi * x)


def bias_variance(make_model, n_sets=500, n=30, sigma=0.3, seed=0):
    rng = np.random.default_rng(seed)
    x_grid = np.linspace(0, 1, 101)
    preds, test_mse = [], []
    for _ in range(n_sets):
        x = rng.uniform(0, 1, n)
        y = f(x) + rng.normal(0, sigma, n)
        p = make_model().fit(x.reshape(-1, 1), y).predict(x_grid.reshape(-1, 1))
        y_new = f(x_grid) + rng.normal(0, sigma, x_grid.size)
        preds.append(p)
        test_mse.append(np.mean((y_new - p) ** 2))
    preds = np.array(preds)
    return (np.mean((preds.mean(axis=0) - f(x_grid)) ** 2), np.mean(preds.var(axis=0)),
            np.mean(test_mse))


poly = lambda d: make_pipeline(PolynomialFeatures(d), StandardScaler(), LinearRegression())
knobs = [
    ('polynomial degree', list(range(0, 13)), lambda d: poly(d), 'degree  (more terms →)'),
    ('k-NN', [30, 20, 15, 10, 7, 5, 3, 2, 1], lambda k: KNeighborsRegressor(n_neighbors=k),
     'k  (fewer neighbors →)'),
    ('regression tree', [1, 2, 3, 4, 5, 6, 8, None],
     lambda d: DecisionTreeRegressor(max_depth=d, random_state=0), 'max_depth  (deeper →)'),
]

fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.5))
for ax, (title, values, make, xlabel) in zip(axes, knobs):
    rows = np.array([bias_variance(lambda: make(v)) for v in values])
    pos = np.arange(len(values))
    total = rows[:, 0] + rows[:, 1] + 0.09
    ax.plot(pos, rows[:, 0], 'o-', color='#f0883e', ms=3.5, label='bias²')
    ax.plot(pos, rows[:, 1], 'o-', color='#58a6ff', ms=3.5, label='variance')
    ax.axhline(0.09, color='gray', ls=':', lw=1.2, label='noise σ² = 0.09')
    ax.plot(pos, total, '-', color='#3fb950', lw=2.2, label='bias² + variance + noise')
    ax.plot(pos, rows[:, 2], 'x', color='#ff7b72', ms=6, mew=1.5, label='measured test MSE')
    best = int(np.argmin(total))
    ax.annotate(f'lowest: {values[best]}\nsum {total[best]:.3f}', xy=(pos[best], total[best]),
                xytext=(pos[best], 3.0 if title == 'polynomial degree' else 0.62), ha='center', fontsize=7.5,
                arrowprops=dict(arrowstyle='->', color='gray'))
    ax.set_yscale('log')
    ax.set_ylim(1e-4, 60 if title == 'polynomial degree' else 1.5)
    ax.set_xticks(pos, [str(v) if v is not None else 'none' for v in values], fontsize=8)
    ax.set_xlabel(xlabel, fontsize=8.5)
    ax.set_title(title, fontsize=10)
    ax.tick_params(axis='y', labelsize=8)
axes[0].set_ylabel('error (log scale)', fontsize=9)
axes[0].text(0.3, 1.6e-4, '← underfits', fontsize=7.5, color='gray')
axes[0].text(12.2, 1.6e-4, 'overfits →', fontsize=7.5, color='gray', ha='right')
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(handles, labels, loc='lower center', ncol=5, fontsize=8, frameon=False)
plt.tight_layout(rect=(0, 0.07, 1, 1))

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

plt.close()