import time
import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from scipy.cluster.hierarchy import linkage
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs

# Left: wall-clock time of scipy's linkage (Ward and single) and of k-means
# (k = 5, one initialization) as n doubles, 5-D blobs; best of 3 runs.
# Right: memory for the n(n-1)/2 pairwise distances in float64, the input
# every standard agglomerative algorithm needs.
ns = np.array([1000, 2000, 4000, 8000, 16000])
X_all, _ = make_blobs(n_samples=ns.max(), n_features=5, centers=5, random_state=0)

def best_of(f, reps=3):
    t = []
    for _ in range(reps):
        t0 = time.perf_counter()
        f()
        t.append(time.perf_counter() - t0)
    return min(t)

times = {'Ward': [], 'single': [], 'k-means': []}
for n in ns:
    X = X_all[:n]
    times['Ward'].append(best_of(lambda: linkage(X, 'ward')))
    times['single'].append(best_of(lambda: linkage(X, 'single')))
    times['k-means'].append(best_of(lambda: KMeans(5, n_init=1, random_state=0).fit(X)))

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, (a1, a2) = plt.subplots(1, 2, figsize=(10, 3.9))

for (name, t), c in zip(times.items(), (RED, BLUE, GREEN)):
    t = np.array(t)
    slope = np.polyfit(np.log(ns), np.log(t), 1)[0]       # time ~ n ** slope
    a1.loglog(ns, t, 'o-', color=c, label=f'{name}: {t[-1]:.3g} s at {ns[-1]:,} rows, time ∝ $n^{{{slope:.1f}}}$')
a1.set_xlabel('number of rows n', fontsize=9)
a1.set_ylabel('seconds (log scale)', fontsize=9)
a1.set_title('Time: the trees grow like $n^2$ or faster', fontsize=9.5, loc='left')
a1.legend(fontsize=8.5, frameon=False, loc='upper left', labelcolor='linecolor')
a1.set_ylim(top=max(max(v) for v in times.values()) * 40)
a1.set_xticks(ns, [f'{n:,}' for n in ns])
a1.minorticks_off()
a1.tick_params(labelsize=8)

big = np.logspace(3, 6, 60)
mem = big * (big - 1) / 2 * 8 / 1e9
a2.loglog(big, mem, color=RED)
for n in (10_000, 100_000, 1_000_000):
    gb = n * (n - 1) / 2 * 8 / 1e9
    a2.plot(n, gb, 'o', color=RED)
    a2.annotate(f'{n:,} rows: {gb:,.1f} GB' if gb < 100 else f'{n:,} rows: {gb:,.0f} GB', (n, gb),
                xytext=(-8, 4 if n < 1_000_000 else -12), textcoords='offset points', ha='right', fontsize=8.5, color=RED)
a2.axhline(16, color=GRAY, ls='--', lw=1)
a2.text(1.1e3, 18, 'a 16 GB laptop', fontsize=8.5, va='bottom')
a2.set_xlabel('number of rows n', fontsize=9)
a2.set_ylabel('GB for all pairwise distances', fontsize=9)
a2.set_title('Memory: n(n−1)/2 distances × 8 bytes', fontsize=9.5, loc='left')
a2.tick_params(labelsize=8)
plt.tight_layout()

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

plt.close()