import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.cluster import HDBSCAN, KMeans
from sklearn.datasets import load_wine
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

# The 178 wines, standardized (13 measurements), projected on their first two
# principal components. The clusterings are computed in 13-D; the plane is only
# a window onto them. Left: k-means, k = 3. Right: HDBSCAN (min_cluster_size=10),
# gray = noise.
df = load_wine(as_frame=True).data
Z = StandardScaler().fit_transform(df)
pca = PCA(n_components=2).fit(Z)
A = pca.transform(Z)
km = KMeans(n_clusters=3, n_init=10, random_state=0).fit(Z)
hd = HDBSCAN(min_cluster_size=10).fit_predict(Z)
PAL = np.array(['#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.8), sharex=True, sharey=True)
a1.scatter(A[:, 0], A[:, 1], s=16, c=PAL[km.labels_], alpha=0.85)
C = pca.transform(km.cluster_centers_)
a1.scatter(C[:, 0], C[:, 1], marker='X', s=200, c=PAL, edgecolor='gray', lw=1.5)
a1.set_title('k-means, k = 3 (fit in 13-D)', fontsize=9.5)
noise = hd == -1
a2.scatter(A[noise, 0], A[noise, 1], s=16, color='gray', alpha=0.5, label=f'noise ({noise.sum()})')
a2.scatter(A[~noise, 0], A[~noise, 1], s=16, c=PAL[hd[~noise]], alpha=0.85)
a2.set_title(f'HDBSCAN: {hd.max() + 1} clusters, {noise.sum()} of 178 wines are noise', fontsize=9.5)
a2.legend(fontsize=8, frameon=False, loc='lower left')
ev = pca.explained_variance_ratio_ * 100
for ax in (a1, a2):
    ax.set_xlabel(f'PC1 ({ev[0]:.1f}% of the variance)', fontsize=8.5)
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.2)
a1.set_ylabel(f'PC2 ({ev[1]:.1f}%)', fontsize=8.5)
plt.tight_layout()

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

plt.close()