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)          # 1,797 images of 8x8 = 64 pixels, values 0-16
pca = PCA().fit(X)                           # pixels share one unit: no standardization
evr = pca.explained_variance_ratio_
cum = np.cumsum(evr)
k = np.arange(1, len(evr) + 1)

fig, ax = plt.subplots(figsize=(8.5, 3.8))
ax.bar(k, evr * 100, color='#58a6ff', alpha=0.8, label='variance of each component (scree)')
ax.set_xlabel('number of components k', fontsize=9)
ax.set_ylabel('% of the variance, each component', fontsize=9, color='#58a6ff')
ax.set_xlim(0, 65)
ax.tick_params(labelsize=8)
ax2 = ax.twinx()
ax2.plot(k, cum * 100, color='#f0883e', lw=2, label='cumulative')
ax2.set_ylabel('cumulative % of the variance', fontsize=9, color='#f0883e')
ax2.set_ylim(0, 103)
ax2.tick_params(labelsize=8)
for t, c in [(0.90, '#3fb950'), (0.95, '#ff7b72')]:
    kt = int(np.searchsorted(cum, t) + 1)
    ax2.axhline(t * 100, color=c, ls='--', lw=1)
    ax2.plot([kt, kt], [0, cum[kt - 1] * 100], color=c, ls=':', lw=1.2)
    ax2.text(kt + 0.8, 38 if t < 0.93 else 28, f'{t:.0%} of the variance → k = {kt}', color=c, fontsize=9)
h1, l1 = ax.get_legend_handles_labels()
h2, l2 = ax2.get_legend_handles_labels()
ax.legend(h1 + h2, l1 + l2, fontsize=8, loc='center right', framealpha=0)
ax.set_title(f'Digits, 64 pixels: the first 2 components keep only {cum[1]:.1%} of the variance',
             fontsize=10, loc='left')
plt.tight_layout()

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

plt.close()