import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from scipy.cluster.hierarchy import dendrogram, fcluster, linkage

# The worked example of the lesson: five points on a line, the gaps between
# neighbors grow by one (3, 4, 5, 6). Top: the points. Bottom: the tree that
# each linkage builds from the same distance matrix, with every merge height.
x = np.array([0, 3, 7, 12, 18])
names = list('ABCDE')
X = x.reshape(-1, 1).astype(float)

GRAY = '#8b949e'
COLORS = {'single': '#58a6ff', 'complete': '#f0883e', 'average': '#3fb950', 'ward': '#ff7b72'}
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})

fig = plt.figure(figsize=(10, 5.2))
gs = fig.add_gridspec(2, 4, height_ratios=[1, 3.2], hspace=0.55, wspace=0.28)

# --- the five points on a line
ax = fig.add_subplot(gs[0, :])
ax.axhline(0, color=GRAY, lw=1, zorder=1)
ax.scatter(x, np.zeros_like(x), s=70, color=GRAY, zorder=3)
for xi, n in zip(x, names):
    ax.annotate(f'{n}\n{xi}', (xi, 0), xytext=(0, 9), textcoords='offset points',
                ha='center', va='bottom', fontsize=9)
for a, b in zip(x[:-1], x[1:]):
    ax.annotate('', xy=(b - 0.15, -0.55), xytext=(a + 0.15, -0.55),
                arrowprops=dict(arrowstyle='<->', color=GRAY, lw=0.8))
    ax.text((a + b) / 2, -0.75, f'gap {b - a}', ha='center', va='top', fontsize=8.5)
ax.set_xlim(-1, 19)
ax.set_ylim(-1.6, 1.6)
ax.axis('off')
ax.set_title('Five points on a line: the gaps between neighbors grow by one', fontsize=10, loc='left')

# --- one dendrogram per linkage, same y scale
for j, method in enumerate(COLORS):
    ax = fig.add_subplot(gs[1, j])
    Z = linkage(X, method)
    color = COLORS[method]
    d = dendrogram(Z, labels=names, ax=ax, link_color_func=lambda k: color)
    # the y ticks are exactly the merge heights, with a faint guide to each bar
    heights = Z[:, 2]
    for xs, ys in zip(d['icoord'], d['dcoord']):
        ax.plot([0, xs[1]], [ys[1], ys[1]], ls=':', lw=0.8, color=color, alpha=0.6)
    ax.set_yticks(heights, [f'{h:.2f}'.rstrip('0').rstrip('.') for h in heights])
    labels = fcluster(Z, 2, criterion='maxclust')
    groups = sorted((''.join(n for n, l in zip(names, labels) if l == c) for c in set(labels)),
                    key=lambda g: g[0])
    ax.set_title(f'{method}\n2 clusters: {{{groups[0]}}} {{{groups[1]}}}', fontsize=9.5, color=color)
    ax.set_ylim(0, 19.5)
    ax.tick_params(labelsize=8.5)
    for t in ax.get_yticklabels():
        t.set_color(color)
    for s in ('top', 'right'):
        ax.spines[s].set_visible(False)
    if j == 0:
        ax.set_ylabel('merge height (a distance)', fontsize=9)

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

plt.close()