import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from scipy.cluster.hierarchy import dendrogram, leaves_list, linkage
from sklearn.datasets import fetch_openml
from sklearn.preprocessing import MinMaxScaler

# The lab's final picture: the 101 animals of the UCI Zoo table, average
# linkage on the scaled features (legs / 8, yes/no as 1/0). Left: the tree,
# with the cuts at 4 and 7 clusters; names colored by the zoologists' class.
# Right: the table itself, rows in leaf order, features clustered the same way.
zoo = fetch_openml(data_id=42667, as_frame=True).frame
animals, kind = zoo['animal'].to_numpy(), zoo['type'].astype(str).to_numpy()
X = zoo.drop(columns=['animal', 'type'])
for col in X.columns.drop('legs'):
    X[col] = (X[col].astype(str) == 'True').astype(int)
Xs = MinMaxScaler().fit_transform(X)
Z = linkage(Xs, 'average')
h = Z[::-1, 2]
rows, cols = leaves_list(Z), leaves_list(linkage(Xs.T, 'average'))

GRAY = '#8b949e'
CLASS = {'mammal': '#f0883e', 'bird': '#58a6ff', 'fish': '#39c5cf', 'reptile': '#3fb950',
         'amphibian': '#bc8cff', 'insect': '#d29922', 'invertebrate': '#ff7b72'}
plt.rcParams.update({'text.color': GRAY, 'axes.labelcolor': GRAY, 'axes.edgecolor': GRAY,
                     'xtick.color': GRAY, 'ytick.color': GRAY})
fig = plt.figure(figsize=(8.6, 14.5))
gs = fig.add_gridspec(1, 2, width_ratios=[1.35, 1], wspace=0.34)

ax = fig.add_subplot(gs[0, 0])
dendrogram(Z, labels=animals, orientation='left', ax=ax, leaf_font_size=7.5,
           link_color_func=lambda k: GRAY)
ax.invert_yaxis()                               # first leaf at the top, like the table
for lab in ax.get_yticklabels():
    lab.set_color(CLASS[kind[list(animals).index(lab.get_text())]])
# a leaf name can appear twice (two frogs); color them both by their class
for k, c in [(4, '#58a6ff'), (7, '#f0883e')]:
    cut = (h[k - 1] + h[k - 2]) / 2
    ax.axvline(cut, color=c, ls='--', lw=1.2)
    ax.text(cut, -8, f'{k} clusters', color=c, ha='center', fontsize=8.5, rotation=90, va='bottom')
ax.set_xlabel('merge height (average distance between the two groups)', fontsize=8.5)
ax.tick_params(axis='x', labelsize=8)
ax.set_ylim(len(animals) * 10, -40)
for s in ('top', 'right', 'left'):
    ax.spines[s].set_visible(False)

axh = fig.add_subplot(gs[0, 1])
axh.imshow(Xs[rows][:, cols], aspect='auto', cmap='Blues', vmin=-0.15, vmax=1,
           interpolation='nearest', extent=(-0.5, len(cols) - 0.5, len(rows) - 0.5, -0.5))
axh.set_ylim(len(rows) - 0.5, -4.5)
axh.set_yticks([])
axh.set_xticks(range(len(cols)), X.columns[cols], rotation=90, fontsize=8)
axh.xaxis.set_ticks_position('top')
axh.set_title('dark = yes (legs: darker = more)', fontsize=8.5, pad=58)

handles = [plt.Line2D([], [], ls='', marker='s', color=c, label=n) for n, c in CLASS.items()]
fig.legend(handles=handles, loc='lower center', ncol=7, fontsize=8, frameon=False,
           labelcolor='linecolor', handletextpad=0.2, columnspacing=0.9)
plt.subplots_adjust(top=0.93, bottom=0.05, left=0.05, right=0.97)

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

plt.close()