import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from numpy.polynomial import Polynomial

# The fundamental lab's world: 15 training points drawn around a smooth curve, and 1,000 new
# points from the same process that the model never sees. Three polynomials, of degree 1, 3 and
# 12, fitted to the 15 points (steps 1 and 2 of the lab: same seeds, same numbers).
def true_curve(x):
    return 0.9 * np.sin(2.2 * np.pi * x) + 0.6 * x

def draw(n, seed):
    rng = np.random.default_rng(seed)
    x = rng.uniform(0, 1, n)
    y = true_curve(x) + rng.normal(0, 0.3, n)
    return x, y

x_train, y_train = draw(15, seed=0)
x_new, y_new = draw(1000, seed=1)
mse = lambda m, x, y: np.mean((y - m(x)) ** 2)

BLUE, ORANGE, GREEN, RED, YELLOW = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72', '#d29922'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 3, figsize=(10.5, 4.1), sharey=True)
grid = np.linspace(0, 1, 400)
verdict = {1: ('underfitting', YELLOW), 3: ('a good balance', GREEN), 12: ('overfitting', RED)}
for ax, degree in zip(axes, [1, 3, 12]):
    m = Polynomial.fit(x_train, y_train, degree)
    ax.scatter(x_new[:150], y_new[:150], s=9, facecolors='none', edgecolors=BLUE, alpha=0.45, lw=0.8,
               label='new points (150 of 1,000)')
    ax.plot(grid, true_curve(grid), '--', color='gray', lw=1.2, label='true curve')
    ax.plot(grid, m(grid), color=verdict[degree][1], lw=2.2, label='fitted model')
    ax.scatter(x_train, y_train, s=26, color=ORANGE, zorder=5, label='15 training points')
    ax.set_ylim(-1.8, 2.2)
    ax.set_xlim(0, 1)
    ax.set_title(f'degree {degree}: {verdict[degree][0]}', fontsize=10, color=verdict[degree][1])
    ax.set_xlabel(f'x\n\nerror (MSE) on the training points: {mse(m, x_train, y_train):.3f}\n'
                  f'error (MSE) on the new points: {mse(m, x_new, y_new):.3f}', fontsize=8.5)
    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()