import warnings
import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import make_moons
from sklearn.model_selection import StratifiedKFold, cross_validate
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

warnings.filterwarnings('ignore')
plt.rcParams.update({'text.color': '#808080', 'axes.labelcolor': '#808080', 'axes.edgecolor': '#999999',
                     'xtick.color': '#808080', 'ytick.color': '#808080'})

# The lesson's moons, solver and folds: width 1, 3, 10, 100 with alpha = 1e-4, then 100 units with alpha = 1
X, y = make_moons(n_samples=300, noise=0.3, random_state=1)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
settings = [(1, 1e-4), (3, 1e-4), (10, 1e-4), (100, 1e-4), (100, 1.0)]

pad = 0.5
x0, x1 = X[:, 0].min() - pad, X[:, 0].max() + pad
y0, y1 = X[:, 1].min() - pad, X[:, 1].max() + pad
xx, yy = np.meshgrid(np.linspace(x0, x1, 250), np.linspace(y0, y1, 250))

fig, axes = plt.subplots(1, 5, figsize=(10, 2.9))
for ax, (units, alpha) in zip(axes, settings):
    net = make_pipeline(StandardScaler(), MLPClassifier(hidden_layer_sizes=(units,), solver='lbfgs',
                                                        alpha=alpha, max_iter=5000, random_state=0))
    r = cross_validate(net, X, y, cv=cv, return_train_score=True)
    net.fit(X, y)
    P = net.predict_proba(np.c_[xx.ravel(), yy.ravel()])[:, 1].reshape(xx.shape)
    ax.contourf(xx, yy, P > 0.5, levels=[-0.5, 0.5, 1.5], colors=['#f0883e', '#58a6ff'], alpha=0.14).set_rasterized(True)
    ax.contour(xx, yy, P, levels=[0.5], colors='gray', linewidths=1.3)
    ax.scatter(X[y == 0, 0], X[y == 0, 1], s=6, color='#f0883e')
    ax.scatter(X[y == 1, 0], X[y == 1, 1], s=6, color='#58a6ff')
    ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
    ax.set_xticks([]); ax.set_yticks([])
    ax.set_title(f"{units} unit{'s' if units > 1 else ''}, α = {alpha:g}\n"
                 f"train {r['train_score'].mean():.3f} · CV {r['test_score'].mean():.3f}", fontsize=8.5)
fig.suptitle('One hidden layer (ReLU, lbfgs) on 300 noisy moons: more units bend the boundary; '
             'α straightens it again', fontsize=9)
plt.tight_layout()

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

plt.close()