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'})

# Three tight groups on a line: blue at 0, orange at 10, green at 100 (same data as the lesson's code)
rng = np.random.default_rng(42)
sizes = np.vstack([rng.normal(0, 1, (100, 2)), rng.normal(0, 10, (100, 2)) + [60, 0]])  # keeps the seed in step
dists = np.vstack([rng.normal(0, 1, (50, 2)) + c for c in [[0, 0], [10, 0], [100, 0]]])
g = np.repeat([0, 1, 2], 50)
pal = np.array(['#58a6ff', '#f0883e', '#3fb950'])


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


fig, axes = plt.subplots(1, 5, figsize=(10.5, 2.9))
axes[0].scatter(dists[:, 0], dists[:, 1], c=pal[g], s=6)
axes[0].set_aspect('equal', adjustable='datalim')
axes[0].set_title(f'original: ratio {dist_ratio(dists):.1f}', fontsize=9, color='gray', loc='left')
for ax, p in zip(axes[1:], [2, 5, 30, 100]):
    E = TSNE(perplexity=p, random_state=0).fit_transform(dists)
    ax.scatter(E[:, 0], E[:, 1], c=pal[g], s=6)
    ax.set_title(f'perplexity {p}: ratio {dist_ratio(E):.1f}', fontsize=9, color='#ff7b72', loc='left')
for ax in axes:
    ax.set_xticks([])
    ax.set_yticks([])
fig.suptitle('Ratio = distance(blue, green) / distance(blue, orange). Same data, same seed; only the perplexity changes.',
             fontsize=9.5, color='gray')
plt.tight_layout()

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

plt.close()