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

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

# Expected learning curves, computed from the known truth: for each training size n,
# 300 training sets from sin(2*pi*x) + noise (sd 0.3); training MSE on the set itself,
# test MSE on 101 fresh points. The median over the 300 sets is drawn (the mean at small n
# is dominated by a few exploding fits), with the band between the 25th and 75th percentiles.
f = lambda x: np.sin(2 * np.pi * x)
sizes = [12, 15, 20, 30, 50, 80, 120, 200, 400, 800]
x_grid = np.linspace(0, 1, 101)
rng = np.random.default_rng(0)


def curves(degree):
    out = []
    for n in sizes:
        tr, te = [], []
        for _ in range(300):
            x = rng.uniform(0, 1, n)
            y = f(x) + rng.normal(0, 0.3, n)
            m = make_pipeline(PolynomialFeatures(degree), StandardScaler(), LinearRegression())
            m.fit(x.reshape(-1, 1), y)
            tr.append(np.mean((m.predict(x.reshape(-1, 1)) - y) ** 2))
            y_new = f(x_grid) + rng.normal(0, 0.3, x_grid.size)
            te.append(np.mean((m.predict(x_grid.reshape(-1, 1)) - y_new) ** 2))
        out.append([np.percentile(tr, [25, 50, 75]), np.percentile(te, [25, 50, 75])])
    return np.array(out)                      # sizes x (train, test) x (p25, p50, p75)


fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.5), sharey=True)
for ax, d, title in [(axes[0], 1, 'degree 1: high bias'), (axes[1], 9, 'degree 9: high variance')]:
    c = curves(d)
    for k, color, label in [(0, '#f0883e', 'training error'), (1, '#58a6ff', 'test error')]:
        ax.fill_between(sizes, c[:, k, 0], c[:, k, 2], color=color, alpha=0.18, lw=0)
        ax.plot(sizes, c[:, k, 1], 'o-', color=color, ms=3.5, lw=2, label=label)
    ax.axhline(0.09, color='gray', ls=':', lw=1.2)
    ax.text(800, 0.082, 'noise floor σ² = 0.09', fontsize=7.5, color='gray', ha='right', va='top')
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_xticks([12, 20, 50, 200, 800], ['12', '20', '50', '200', '800'], fontsize=8)
    ax.set_xlabel('training-set size n')
    ax.set_title(title, fontsize=10)
axes[0].annotate('curves meet near 0.29 = bias² + noise,\nfar above the floor:\nmore data changes nothing', xy=(400, 0.28),
                 xytext=(40, 0.9), fontsize=8, arrowprops=dict(arrowstyle='->', color='gray'))
axes[1].annotate('gap closes as n grows:\nmore data helps', xy=(120, 0.11), xytext=(160, 0.9),
                 fontsize=8, arrowprops=dict(arrowstyle='->', color='gray'))
axes[0].set_ylabel('MSE (median, log scale)')
axes[0].set_ylim(0.01, 8)
axes[0].legend(fontsize=8, loc='lower right', frameon=False)
plt.tight_layout()

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

plt.close()