import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import RepeatedKFold, validation_curve
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'})

# The lesson's validation curve (one sample of 30 points, seed 1), and the same
# curve for 20 other samples of 30 points from the same distribution.
f = lambda x: np.sin(2 * np.pi * x)
degrees = np.arange(1, 11)
model = make_pipeline(PolynomialFeatures(), StandardScaler(), LinearRegression())
rcv = RepeatedKFold(n_splits=5, n_repeats=20, random_state=0)


def curve(seed):
    rng = np.random.default_rng(seed)
    x = rng.uniform(0, 1, 30)
    y = f(x) + rng.normal(0, 0.3, 30)
    tr, va = validation_curve(model, x.reshape(-1, 1), y, param_name='polynomialfeatures__degree',
                              param_range=degrees, cv=rcv, scoring='neg_mean_squared_error')
    return -tr.mean(axis=1), -va.mean(axis=1)


fig, ax = plt.subplots(figsize=(7.2, 3.6))
others = [curve(s)[1] for s in range(100, 120)]
for i, va in enumerate(others):
    ax.plot(degrees, va, color='#58a6ff', lw=0.7, alpha=0.3,
            label='validation, 20 other samples of 30' if i == 0 else None)
tr, va = curve(1)
ax.plot(degrees, tr, 'o-', color='#f0883e', lw=2, ms=4, label='training (this sample)')
ax.plot(degrees, va, 'o-', color='#58a6ff', lw=2, ms=4, label='validation (this sample)')
ax.axhline(0.09, color='gray', ls=':', lw=1.2)
ax.text(10.1, 0.082, 'noise σ² = 0.09', fontsize=7.5, color='gray', ha='right', va='top')
best = int(np.argmin(va))
ax.annotate(f'bottom: degree {degrees[best]}, validation MSE {va[best]:.3f}',
            xy=(degrees[best], va[best]), xytext=(3.3, 0.021), fontsize=8,
            arrowprops=dict(arrowstyle='->', color='gray'))
ax.text(1.0, 0.06, 'bias side:\nboth curves high,\nsmall gap', fontsize=7.5, color='gray', va='top')
ax.text(10.1, 0.036, 'variance side:\ntraining keeps falling,\nvalidation climbs', fontsize=7.5,
        color='gray', va='top', ha='right')
ax.set_yscale('log')
ax.set_ylim(0.015, 3)
ax.set_xticks(degrees)
ax.set_xlabel('polynomial degree')
ax.set_ylabel('MSE (log scale), 5-fold × 20')
ax.legend(fontsize=7.5, loc='upper left', frameon=False)
ax.set_title('validation_curve on one sample of 30 points', fontsize=10)
plt.tight_layout()

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

plt.close()