import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LinearRegression
from sklearn.neighbors import KNeighborsRegressor

# No free lunch, on two worlds (step 5 of the advanced lab: same seeds, same numbers).
# 30 training points from a straight world (y = 2x + noise) and from a wiggly one (the
# fundamental lab's curve + noise); a straight line and 5 nearest neighbors fitted to each,
# both scored by the mean squared error on 1,000 new points from the same world.
def true_curve(x):
    return 0.9 * np.sin(2.2 * np.pi * x) + 0.6 * x

worlds = {'straight': lambda x: 2 * x, 'wiggly': true_curve}

def draw_world(world, n, seed):
    rng = np.random.default_rng(seed)
    x = rng.uniform(0, 1, n)
    return x.reshape(-1, 1), worlds[world](x) + rng.normal(0, 0.3, n)

BLUE, ORANGE, GREEN, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#bc8cff'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.1))
grid = np.linspace(0, 1, 400).reshape(-1, 1)
for ax, world in zip(axes, worlds):
    X, y = draw_world(world, 30, seed=0)
    X_new, y_new = draw_world(world, 1000, seed=1)
    line = LinearRegression().fit(X, y)
    knn = KNeighborsRegressor(n_neighbors=5).fit(X, y)
    e_line = np.mean((y_new - line.predict(X_new)) ** 2)
    e_knn = np.mean((y_new - knn.predict(X_new)) ** 2)
    ax.plot(grid[:, 0], worlds[world](grid[:, 0]), '--', color='gray', lw=1.2, label='true curve')
    ax.plot(grid[:, 0], line.predict(grid), color=BLUE, lw=2.2, label='straight line')
    ax.plot(grid[:, 0], knn.predict(grid), color=PURPLE, lw=2.2, label='5 nearest neighbors')
    ax.scatter(X[:, 0], y, s=24, color=ORANGE, zorder=5, label='30 training points')
    win = 'straight line' if e_line < e_knn else '5 nearest neighbors'
    ax.set_title(f'{world} world: the {win} wins', fontsize=10, color=BLUE if e_line < e_knn else PURPLE)
    ax.set_xlabel(f'x\n\nerror (MSE) on 1,000 new points: line {e_line:.3f} · neighbors {e_knn:.3f}', fontsize=8.5)
    ax.set_xlim(0, 1)
    ax.set_ylim(-1.6, 2.9)
    ax.spines[['top', 'right']].set_visible(False)
axes[0].set_ylabel('y')
h, l = axes[0].get_legend_handles_labels()
fig.legend(h, l, loc='upper center', ncol=4, frameon=False, fontsize=8)
plt.tight_layout(rect=(0, 0, 1, 0.93))

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

plt.close()