import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from scipy.cluster.hierarchy import dendrogram, fcluster, linkage
from sklearn.datasets import make_blobs
from sklearn.metrics import silhouette_score

# 400 points in four blobs that come in two pairs: structure at two scales.
# Ward linkage. Top: the tree (last 16 merges), the range of cut heights that
# gives each number of clusters, and the silhouette of each cut.
# Bottom: the two cuts the gaps point to.
X, y = make_blobs(n_samples=400, centers=[[0, 0], [0, 3.5], [9, 6], [12, 6]],
                  cluster_std=0.8, random_state=0)
Z = linkage(X, 'ward')
h = Z[::-1, 2]                                  # h[0] merges 2 -> 1 clusters, h[1] 3 -> 2, ...
ks = np.arange(2, 10)
gap = h[ks - 2] - h[ks - 1]                     # range of heights where the cut gives k clusters
sil = [silhouette_score(X, fcluster(Z, k, criterion='maxclust')) for k in ks]

GRAY = '#8b949e'
BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})
fig = plt.figure(figsize=(10, 6.0))
gs = fig.add_gridspec(2, 3, height_ratios=[1.1, 1], hspace=0.45, wspace=0.32)

ax = fig.add_subplot(gs[0, 0])
dendrogram(Z, truncate_mode='lastp', p=16, ax=ax, link_color_func=lambda k: GRAY,
           no_labels=True)
cut2, cut4 = (h[0] + h[1]) / 2, (h[2] + h[3]) / 2
ax.axhline(cut2, color=ORANGE, ls='--', lw=1.3, label='cut → 2 clusters')
ax.axhline(cut4, color=BLUE, ls='--', lw=1.3, label='cut → 4 clusters')
ax.set_ylim(0, h[0] * 1.38)
ax.legend(loc='upper right', fontsize=8, frameon=False, labelcolor='linecolor', ncol=1)
ax.set_ylabel('merge height', fontsize=9)
ax.set_title('Ward tree (last 16 merges)', fontsize=9.5, loc='left')
ax.tick_params(labelsize=8)

ax = fig.add_subplot(gs[0, 1])
colors = [ORANGE if k == 2 else BLUE if k == 4 else GRAY for k in ks]
ax.bar(ks, gap, color=colors)
for k, g in zip(ks, gap):
    ax.text(k, g + 2, f'{g:.1f}', ha='center', va='bottom', fontsize=8)
ax.set_xlabel('number of clusters k', fontsize=9)
ax.set_ylabel('height range that gives k', fontsize=9)
ax.set_title('Gap in merge heights', fontsize=9.5, loc='left')
ax.set_xticks(ks)
ax.set_ylim(0, gap.max() * 1.15)
ax.tick_params(labelsize=8)

ax = fig.add_subplot(gs[0, 2])
ax.plot(ks, sil, 'o-', color=GREEN)
for k, s in zip(ks, sil):
    if k in (2, 4):
        ax.annotate(f'{s:.3f}', (k, s), xytext=(6, 2), textcoords='offset points', fontsize=8.5,
                    color=ORANGE if k == 2 else BLUE)
ax.set_xlabel('number of clusters k', fontsize=9)
ax.set_ylabel('mean silhouette', fontsize=9)
ax.set_title('Silhouette of each cut', fontsize=9.5, loc='left')
ax.set_ylim(top=0.86)
ax.set_xticks(ks)
ax.tick_params(labelsize=8)

PAL = np.array([BLUE, ORANGE, GREEN, RED])
for j, k in enumerate((2, 4)):
    ax = fig.add_subplot(gs[1, j])
    lab = fcluster(Z, k, criterion='maxclust') - 1
    ax.scatter(X[:, 0], X[:, 1], c=PAL[lab], s=7, linewidths=0, rasterized=True)
    ax.set_title(f'k = {k}: silhouette {sil[k - 2]:.3f}', fontsize=9.5, loc='left',
                 color=ORANGE if k == 2 else BLUE)
    ax.set_xticks([])
    ax.set_yticks([])

ax = fig.add_subplot(gs[1, 2])
ax.axis('off')
ax.text(0, 0.85, 'Two cuts, both defensible:', fontsize=9.5, va='top')
ax.text(0, 0.62, f'k = 2 survives a height range\nof {gap[0]:.1f}: two far-apart pairs', fontsize=9,
        va='top', color=ORANGE)
ax.text(0, 0.28, f'k = 4 survives {gap[2]:.1f}: inside each\npair, two blobs', fontsize=9,
        va='top', color=BLUE)

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

plt.close()