import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.cluster import AgglomerativeClustering
from sklearn.datasets import make_blobs, make_moons
from sklearn.metrics import adjusted_rand_score

# Four 2-D datasets with known groups (300 points each), four linkages, the
# true number of clusters given to every method. Color = cluster found;
# the panel title is the ARI against the true groups (1 = perfect, 0 = chance).
rng = np.random.default_rng(0)
blobs = make_blobs(n_samples=300, centers=3, cluster_std=1.0, random_state=1)
a, b = rng.normal([0, 0], 0.7, (140, 2)), rng.normal([6, 0], 0.7, (140, 2))
t = np.linspace(1.0, 5.0, 20)                                     # a thin line of 20 points
bridge = (np.vstack([a, b, np.c_[t, rng.normal(0, 0.05, 20)]]),
          np.r_[np.zeros(140), np.ones(140), (t > 3).astype(int)])
moons = make_moons(n_samples=300, noise=0.06, random_state=0)
rng = np.random.default_rng(0)
strips = (np.vstack([np.c_[rng.uniform(0, 8, 100), rng.normal(1.5 * j, 0.12, 100)] for j in range(3)]),
          np.repeat([0, 1, 2], 100))
data = {'blobs': blobs, 'two blobs + bridge': bridge, 'moons': moons, 'parallel strips': strips}
methods = ['single', 'complete', 'average', 'ward']

GRAY = '#8b949e'
PALETTE = np.array(['#58a6ff', '#f0883e', '#3fb950'])
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})
fig, axes = plt.subplots(4, 4, figsize=(9.5, 8.6))
for i, (name, (X, y)) in enumerate(data.items()):
    k = len(np.unique(y))
    for j, method in enumerate(methods):
        ax = axes[i, j]
        lab = AgglomerativeClustering(n_clusters=k, linkage=method).fit_predict(X)
        # color clusters by size, so the largest is always blue
        order = np.argsort(-np.bincount(lab))
        rank = np.empty_like(order)
        rank[order] = np.arange(k)
        ax.scatter(X[:, 0], X[:, 1], c=PALETTE[rank[lab]], s=6, linewidths=0, rasterized=True)
        tiny = np.bincount(lab)[lab] < 5                          # circle clusters of 1-4 points
        ax.scatter(X[tiny, 0], X[tiny, 1], s=90, facecolors='none', edgecolors=PALETTE[rank[lab[tiny]]],
                   linewidths=1.3)
        score = adjusted_rand_score(y, lab)
        good = score > 0.9
        ax.set_title(f'ARI {score:.2f}', fontsize=9.5, color='#3fb950' if good else '#ff7b72',
                     fontweight='bold' if good else 'normal')
        ax.set_xticks([])
        ax.set_yticks([])
        if i == 0:
            ax.text(0.5, 1.25, method, transform=ax.transAxes, ha='center', fontsize=11)
        if j == 0:
            ax.set_ylabel(name, fontsize=10)
fig.suptitle('Same number of clusters, four linkages: each one wins where its idea of a cluster holds',
             fontsize=10.5, y=0.995)
plt.tight_layout(rect=(0, 0, 1, 0.975))

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

plt.close()