import os

# 限制 BLAS/MKL/OpenMP 只用 1 个线程
os.environ["OMP_NUM_THREADS"]         = "1"   # OpenMP
os.environ["OPENBLAS_NUM_THREADS"]    = "1"   # OpenBLAS
os.environ["MKL_NUM_THREADS"]         = "1"   # Intel MKL
os.environ["VECLIB_MAXIMUM_THREADS"]  = "1"   # Mac OS X Accelerate
os.environ["NUMEXPR_NUM_THREADS"]     = "1"   # NumExpr

import time

import numpy as np
from sklearn.manifold import TSNE

import torch
torch.set_num_threads(1)
torch.set_num_interop_threads(1)


tsne = TSNE(perplexity=20, n_components=2, verbose=0, init='random', random_state=42, n_jobs=1, min_grad_norm=1e-18,
            n_iter=2000, method='exact')

used_times = []
for i in range(10):
    x = np.random.randn(3000, 512)
    start_time = time.perf_counter()
    z = tsne.fit_transform(x)
    t = time.perf_counter() - start_time
    used_times.append(t)
    print(t)

used_times = np.array(used_times)
torch.save(used_times, './tsne_used_times.save')
print(used_times.mean(), used_times.std())
