import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.cluster import HDBSCAN, KMeans
from sklearn.datasets import make_blobs, make_circles, make_moons
from sklearn.metrics import adjusted_rand_score, silhouette_score

# Five 2-D datasets of 500 points whose true groups are known. Top row:
# k-means told the true number of groups. Bottom row: HDBSCAN, told nothing
# but min_cluster_size=15; gray points are what it calls noise.
Xa, ya = make_blobs(n_samples=500, centers=3, random_state=170)
shapes = {
    'blobs':       make_blobs(n_samples=500, centers=3, cluster_std=1.0, random_state=8),
    'moons':       make_moons(n_samples=500, noise=0.06, random_state=0),
    'circles':     make_circles(n_samples=500, factor=0.5, noise=0.05, random_state=0),
    'anisotropic': (Xa @ np.array([[0.6, -0.6], [-0.4, 0.8]]), ya),
    'densities':   make_blobs(n_samples=[150, 150, 200], centers=[[0, 0], [1.2, 0], [4, 2]],
                              cluster_std=[0.2, 0.2, 1.0], random_state=0),
}

PAL = np.array(['#58a6ff', '#f0883e', '#3fb950', '#ff7b72', '#d2a8ff'])
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(2, 5, figsize=(10.5, 4.6))
for c, (name, (X, y)) in enumerate(shapes.items()):
    k = len(np.unique(y))
    km = KMeans(n_clusters=k, n_init=10, random_state=0).fit_predict(X)
    hd = HDBSCAN(min_cluster_size=15).fit_predict(X)
    for r, (lab, what) in enumerate([(km, 'k-means'), (hd, 'HDBSCAN')]):
        ax = axes[r, c]
        noise = lab == -1
        ax.scatter(X[noise, 0], X[noise, 1], s=5, color='gray', alpha=0.5, rasterized=True)
        ax.scatter(X[~noise, 0], X[~noise, 1], s=5, c=PAL[lab[~noise] % 5], alpha=0.8, rasterized=True)
        ari = adjusted_rand_score(y, lab)
        extra = f' · sil {silhouette_score(X, km):.2f}' if r == 0 else f' · {noise.sum()} noise'
        good = ari > 0.8
        ax.set_title(f'ARI {ari:.2f}{extra}'.replace('-0.00', '0.00'), fontsize=8.5,
                     color='#3fb950' if good else '#ff7b72')
        ax.set_xticks([])
        ax.set_yticks([])
        if r == 0:
            ax.text(0.5, 1.2, name, transform=ax.transAxes, ha='center', fontsize=10)
        if c == 0:
            ax.set_ylabel(what + (f'\n(k given)' if r == 0 else '\n(no k)'), fontsize=9)
fig.suptitle('k-means needs round, similar groups; the silhouette does not warn you when they are not',
             fontsize=10, y=0.99)
plt.tight_layout(rect=(0, 0, 1, 0.97))

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

plt.close()