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.neighbors import KNeighborsClassifier

# The lab's step 7, drawn: a test digit that 3-NN gets right and the first one it gets
# wrong, each followed by its 5 nearest training images
plt.rcParams.update({'text.color': 'gray'})
digits = load_digits()
Xd, yd = digits.data, digits.target
Xd_train, Xd_test, yd_train, yd_test = train_test_split(Xd, yd, test_size=0.25,
                                                        stratify=yd, random_state=0)
knn_d = KNeighborsClassifier(n_neighbors=3).fit(Xd_train, yd_train)
pred = knn_d.predict(Xd_test)
wrong = np.flatnonzero(pred != yd_test)

fig, axes = plt.subplots(2, 6, figsize=(8, 3.3))
for row, i in zip(axes, [0, wrong[0]]):
    dist, idx = knn_d.kneighbors(Xd_test[[i]], n_neighbors=5)
    row[0].imshow(Xd_test[i].reshape(8, 8), cmap='gray_r')
    ok = pred[i] == yd_test[i]
    row[0].set_title(f'test: true {yd_test[i]}\npredicted {pred[i]}', fontsize=9,
                     color='#3fb950' if ok else '#ff7b72')
    for ax, j, d in zip(row[1:], idx[0], dist[0]):
        ax.imshow(Xd_train[j].reshape(8, 8), cmap='gray_r')
        ax.set_title(f'neighbor: {yd_train[j]}\nd = {d:.1f}', fontsize=8.5)
    for ax in row:
        ax.set_xticks([])
        ax.set_yticks([])

plt.tight_layout()

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

plt.close()