import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.manifold import TSNE

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# The same three toy datasets as the code blocks of the lesson
rng = np.random.default_rng(42)
sizes = np.vstack([rng.normal(0, 1, (100, 2)),                  # a tight group (std 1)
                   rng.normal(0, 10, (100, 2)) + [60, 0]])      # a loose group (std 10)
g_sizes = np.repeat([0, 1], 100)
dists = np.vstack([rng.normal(0, 1, (50, 2)) + c for c in [[0, 0], [10, 0], [100, 0]]])
g_dists = np.repeat([0, 1, 2], 50)
noise = rng.random((200, 10))                                   # 10 columns of pure noise


def spread(Z):
    return np.sqrt(((Z - Z.mean(axis=0)) ** 2).sum(axis=1).mean())


def size_ratio(Z):
    return spread(Z[g_sizes == 1]) / spread(Z[g_sizes == 0])


def dist_ratio(Z):
    m = [Z[g_dists == k].mean(axis=0) for k in range(3)]
    return np.linalg.norm(m[0] - m[2]) / np.linalg.norm(m[0] - m[1])


E_sizes = TSNE(perplexity=30, random_state=0).fit_transform(sizes)
E_dists = TSNE(perplexity=30, random_state=0).fit_transform(dists)
E_noise = TSNE(perplexity=2, random_state=0).fit_transform(noise)

pal = np.array(['#58a6ff', '#f0883e', '#3fb950'])
fig, axes = plt.subplots(2, 3, figsize=(10, 6.2))
panels = [
    (sizes, E_sizes, g_sizes,
     f'loose group is {size_ratio(sizes):.1f}× wider', f'loose group is {size_ratio(E_sizes):.1f}× wider'),
    (dists, E_dists, g_dists,
     f'green is {dist_ratio(dists):.1f}× farther than orange', f'green is {dist_ratio(E_dists):.1f}× farther than orange'),
    (noise, E_noise, np.zeros(len(noise), int),
     'uniform noise (first 2 of 10 columns)', 'the same noise, mapped (perplexity 2)'),
]
names = ['1. "bigger group = more spread"', '2. "closer groups = more similar"', '3. "there are groups"']
for j, (X, E, g, t_top, t_bot) in enumerate(panels):
    for i, (Z, t) in enumerate([(X[:, :2], t_top), (E, t_bot)]):
        ax = axes[i, j]
        ax.scatter(Z[:, 0], Z[:, 1], c=pal[g], s=7, alpha=0.8)
        ax.set_xticks([])
        ax.set_yticks([])
        if j < 2 and i == 0:
            ax.set_aspect('equal', adjustable='datalim')
        ax.set_title(t, fontsize=9, color='gray' if i == 0 else '#ff7b72', loc='left')
    axes[0, j].text(0, 1.2, names[j], transform=axes[0, j].transAxes, fontsize=10)
axes[0, 0].set_ylabel('original data', fontsize=9)
axes[1, 0].set_ylabel('t-SNE (perplexity 30 unless noted)', fontsize=9)
plt.tight_layout()

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

plt.close()