# use t-SNE to show the feature distribution

import numpy as np
import matplotlib.pyplot as plt
from sklearn import manifold, datasets
from einops import reduce
import scipy.io
import torch

def plt_tsne(data, label, per):
    data = data.cpu().detach().numpy()
    data = reduce(data, 'b n e -> b e', reduction='mean')
    label = label.cpu().detach().numpy()

    tsne = manifold.TSNE(n_components=2, perplexity=per, init='pca', random_state=166)
    X_tsne = tsne.fit_transform(data)

    x_min, x_max = X_tsne.min(0), X_tsne.max(0)
    X_norm = (X_tsne - x_min) / (x_max - x_min)
    plt.figure()
    for i in range(X_norm.shape[0]):
        plt.scatter(X_norm[i, 0], X_norm[i, 1], color=plt.Set1(label[i]))
        plt.xticks([])
        plt.yticks([])
    # plt.show()
    plt.savefig('./test.png', dpi=600)