import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

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

# Step 6 of the lab: the grid's winner (256 units, alpha = 1e-4), refit without and with early stopping
X, y = load_digits(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=42)
settings = dict(hidden_layer_sizes=(256,), alpha=0.0001, max_iter=1000, random_state=42)
full = make_pipeline(StandardScaler(), MLPClassifier(**settings)).fit(X_train, y_train)[-1]
early = make_pipeline(StandardScaler(), MLPClassifier(**settings, early_stopping=True)).fit(X_train, y_train)[-1]

fig, (a1, a2) = plt.subplots(1, 2, figsize=(9.5, 3.4))
a1.plot(np.arange(1, full.n_iter_ + 1), full.loss_curve_, color='#58a6ff', lw=2,
        label=f'no early stopping: {full.n_iter_} epochs')
a1.plot(np.arange(1, early.n_iter_ + 1), early.loss_curve_, color='#f0883e', lw=2, ls='--',
        label=f'early stopping: {early.n_iter_} epochs')
a1.set_yscale('log')
a1.set_xlabel('epoch', fontsize=9); a1.set_ylabel('training loss (log scale)', fontsize=9)
a1.set_title(f'training loss: {full.loss_curve_[0]:.2f} at epoch 1, {full.loss_curve_[-1]:.4f} at the end',
             fontsize=9)
a1.legend(fontsize=8, framealpha=0)

v = np.array(early.validation_scores_)
best = int(np.argmax(v))
a2.plot(np.arange(1, len(v) + 1), v, 'o-', ms=3, color='#3fb950', lw=1.6)
a2.axvline(best + 1, color='gray', lw=0.8, ls=':')
a2.annotate(f'best {v[best]:.3f} at epoch {best + 1}', xy=(best + 1, v[best]), xytext=(best - 6, 0.72),
            fontsize=8, arrowprops=dict(arrowstyle='->', color='gray'))
a2.text(len(v), v[-1] - 0.035, f'stops at epoch {len(v)}:\nno gain after\nepoch {best + 1}', fontsize=7.5,
        ha='right', va='top')
a2.set_xlabel('epoch', fontsize=9); a2.set_ylabel('accuracy on the 10% held out', fontsize=9)
a2.set_title('early stopping: validation accuracy per epoch', fontsize=9)
for ax in (a1, a2):
    ax.tick_params(labelsize=7.5)
plt.tight_layout()

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

plt.close()