import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import ListedColormap
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

# The lesson's moons: 1000 points, 600 for training, 200 for validation (test untouched)
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
X, y = make_moons(n_samples=1000, noise=0.35, random_state=0)
X_tmp, X_test, y_tmp, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
X_train, X_val, y_train, y_val = train_test_split(X_tmp, y_tmp, test_size=0.25,
                                                  stratify=y_tmp, random_state=0)

xx, yy = np.meshgrid(np.linspace(-2.2, 3.2, 300), np.linspace(-1.8, 2.3, 300))
grid = np.c_[xx.ravel(), yy.ravel()]
cmap_bg = ListedColormap(['#f0883e', '#58a6ff'])

fig, axes = plt.subplots(1, 3, figsize=(10, 3.6))
for ax, k, label in zip(axes, [1, 15, 451], ['memorizes noise', 'follows the moons', 'too smooth']):
    clf = KNeighborsClassifier(n_neighbors=k).fit(X_train, y_train)
    Z = clf.predict(grid).reshape(xx.shape)
    cs = ax.contourf(xx, yy, Z, levels=[-0.5, 0.5, 1.5], cmap=cmap_bg, alpha=0.25)
    cs.set_rasterized(True)
    ax.scatter(X_train[y_train == 0, 0], X_train[y_train == 0, 1], c='#f0883e', s=7, edgecolors='none', rasterized=True)
    ax.scatter(X_train[y_train == 1, 0], X_train[y_train == 1, 1], c='#58a6ff', s=7, edgecolors='none', rasterized=True)
    ax.set_title(f'k = {k} ({label})\ntrain {clf.score(X_train, y_train):.1%} · '
                 f'validation {clf.score(X_val, y_val):.1%}', fontsize=9.5)
    ax.set_xticks([])
    ax.set_yticks([])

plt.tight_layout()

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

plt.close()