import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.cluster import KMeans
from sklearn.datasets import fetch_openml
from sklearn.metrics import adjusted_rand_score, silhouette_score
from sklearn.preprocessing import StandardScaler

# Palmer penguins, 4 standardized measurements, species hidden from k-means.
# Left: the elbow of the inertia. Right: the silhouette (what you can compute
# without labels) against the ARI with the species (what you cannot).
peng = fetch_openml(data_id=42585, as_frame=True).frame
num = ['culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g']
peng = peng.dropna(subset=num)
Z = StandardScaler().fit_transform(peng[num])
ks = np.arange(2, 9)
inertia, sil, ari = [], [], []
for k in ks:
    m = KMeans(n_clusters=k, n_init=10, random_state=0).fit(Z)
    inertia.append(m.inertia_)
    sil.append(silhouette_score(Z, m.labels_))
    ari.append(adjusted_rand_score(peng['species'], m.labels_))
inertia, sil, ari = map(np.array, (inertia, sil, ari))

BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, (a1, a2) = plt.subplots(1, 2, figsize=(9, 3.4))
a1.plot(ks, inertia, 'o-', color=BLUE)
for k in (2, 3):
    i = k - 2
    a1.annotate(f'−{inertia[i] - inertia[i + 1]:.0f}', ((k + 0.5), (inertia[i] + inertia[i + 1]) / 2),
                textcoords='offset points', xytext=(8, 4), fontsize=8.5, color=ORANGE)
a1.scatter([3], [inertia[1]], s=180, facecolor='none', edgecolor=ORANGE, lw=1.6)
a1.set_title('inertia: the elbow is at k = 3', fontsize=9.5)
a1.set_xlabel('k')
a1.set_ylabel('inertia')

a2.plot(ks, sil, 'o-', color=BLUE, label='silhouette (no labels needed)')
a2.plot(ks, ari, 's--', color=GREEN, label='ARI vs species (needs labels)')
for arr, c in [(sil, BLUE), (ari, GREEN)]:
    i = arr.argmax()
    a2.scatter([ks[i]], [arr[i]], s=180, facecolor='none', edgecolor=c, lw=1.6)
    a2.annotate(f'max at k = {ks[i]}', (ks[i], arr[i]), textcoords='offset points',
                xytext=(10, 4), fontsize=8.5, color=c)
a2.set_title('the silhouette peaks at k = 2; there are 3 species', fontsize=9.5)
a2.set_xlabel('k')
a2.set_ylabel('score')
a2.set_ylim(0.25, 0.88)
a2.legend(fontsize=8, frameon=False, loc='upper right')
for ax in (a1, a2):
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.2)
plt.tight_layout()

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

plt.close()