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

# Train and validation accuracy of k-NN on the lesson's moons, for every odd k
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)
ks = np.arange(1, 452, 2)
tr, va = [], []
for k in ks:
    m = KNeighborsClassifier(n_neighbors=k).fit(X_train, y_train)
    tr.append(m.score(X_train, y_train))
    va.append(m.score(X_val, y_val))
tr, va = np.array(tr), np.array(va)
best = int(np.argmax(va))

fig, ax = plt.subplots(figsize=(7.5, 3.8))
ax.plot(ks, tr, color='#f0883e', lw=1.6, label='training accuracy')
ax.plot(ks, va, color='#58a6ff', lw=1.6, label='validation accuracy (200 rows)')
ax.scatter([ks[best]], [va[best]], s=60, facecolors='none', edgecolors='#3fb950', lw=1.6, zorder=5)
ax.annotate(f'best on validation: k = {ks[best]}, {va[best]:.1%}', (ks[best], va[best]),
            xytext=(0.62, 0.95), textcoords='axes fraction', fontsize=8.5,
            arrowprops=dict(arrowstyle='->', color='gray'))
ax.annotate(f'k = 1: training {tr[0]:.0%},\nvalidation {va[0]:.1%}', (1, tr[0]), xytext=(2.2, 0.99),
            fontsize=8.5, va='top', arrowprops=dict(arrowstyle='->', color='gray'))
ax.text(1.4, 0.78, '← small k: complex,\n    high variance', fontsize=8.5)
ax.text(200, 0.72, 'large k: smooth, high bias →', fontsize=8.5, ha='right')
ax.set_xscale('log')
ax.set_xticks([1, 3, 10, 30, 100, 300, 451], ['1', '3', '10', '30', '100', '300', '451'])
ax.set_xlabel('k (number of neighbors), log scale')
ax.set_ylabel('accuracy')
ax.set_ylim(0.66, 1.03)
ax.legend(fontsize=8, loc='lower left', frameon=False)
ax.set_title('Choosing k on a validation split (600 training rows)', fontsize=10)
ax.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()