import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap
from sklearn.datasets import load_digits, make_swiss_roll
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE, trustworthiness

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

# Row 1: handwritten digits, 64 pixels -> 2
X_dig, y_dig = load_digits(return_X_y=True)
# Row 2: a 2-D sheet rolled up in 3-D; t = position along the roll
X_roll, t = make_swiss_roll(n_samples=1000, noise=0.1, random_state=0)
sheet = np.column_stack([t, X_roll[:, 1]])        # the true, unrolled coordinates

plt.rcParams['savefig.dpi'] = 85             # resolution of the rasterized dots
roll_cmap = LinearSegmentedColormap.from_list('roll', ['#58a6ff', '#3fb950', '#f0883e'])
fig, axes = plt.subplots(2, 2, figsize=(9, 7.6))
for i, (X, ref, color, cmap, what) in enumerate([
        (X_dig, X_dig, y_dig, 'tab10', 'digits (color = true digit)'),
        (X_roll, sheet, t, roll_cmap, 'swiss roll (color = position along the roll)')]):
    for j, (name, model) in enumerate([('PCA', PCA(n_components=2)),
                                       ('t-SNE', TSNE(perplexity=30, random_state=0))]):
        E = model.fit_transform(X)
        tw = trustworthiness(ref, E, n_neighbors=12)
        ax = axes[i, j]
        ax.scatter(E[:, 0], E[:, 1], c=color, cmap=cmap, s=5, alpha=0.85, rasterized=True)
        ax.set_title(f'{name} — {what}\ntrustworthiness {tw:.3f}', fontsize=9, loc='left')
        ax.set_xticks([])
        ax.set_yticks([])
        if i == 0:                                 # write each digit at the middle of its points
            for d in range(10):
                mx, my = np.median(E[y_dig == d], axis=0)
                ax.text(mx, my, str(d), fontsize=11, weight='bold', ha='center', va='center',
                        color='white', bbox=dict(boxstyle='circle,pad=0.15', fc='#555555', ec='none', alpha=0.8))
plt.tight_layout(rect=(0, 0.03, 1, 1))
fig.text(0.01, 0.005, 'Trustworthiness (k = 12): 1 means every neighbor on the map was a neighbor in the data '
         '(for the roll: on the unrolled sheet).', fontsize=8, color='gray')

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

plt.close()