import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import load_digits
from sklearn.decomposition import PCA

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

X, y = load_digits(return_X_y=True)
examples = [int(np.flatnonzero(y == d)[0]) for d in [0, 2, 4, 8]]   # the first 0, 2, 4 and 8
ks = [1, 2, 5, 10, 20, 29, 64]

fig, axes = plt.subplots(len(examples), len(ks) + 1, figsize=(9.0, 5.4))
for j, k in enumerate(ks + ['original']):
    if k == 'original':
        R, header = X, 'original\n64 numbers'
    else:
        pca = PCA(n_components=k).fit(X)
        R = pca.inverse_transform(pca.transform(X))
        mse = ((X - R) ** 2).sum(axis=1).mean()
        kept = pca.explained_variance_ratio_.sum()
        header = f'k = {k}\n{kept:.0%} kept\nerror {mse:.0f}'
    for i, idx in enumerate(examples):
        ax = axes[i, j]
        ax.imshow(R[idx].reshape(8, 8), cmap='gray_r', vmin=0, vmax=16)
        ax.set_xticks([])
        ax.set_yticks([])
        for s in ax.spines.values():
            s.set_edgecolor('#58a6ff' if k == 29 else 'gray')
            s.set_linewidth(2 if k == 29 else 0.5)
        if i == 0:
            ax.set_title(header, fontsize=8.5, color='#58a6ff' if k == 29 else 'gray')
fig.suptitle('Rebuilding digits from k principal components (error = mean squared error per image, all 1,797 digits)',
             fontsize=9.5, color='gray')
plt.tight_layout()

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

plt.close()