#
#  tsne_torch.py
#
# Implementation of t-SNE in pytorch. The implementation was tested on pytorch
# > 1.0, and it requires Numpy to read files. In order to plot the results,
# a working installation of matplotlib is required.
#
#
# The example can be run by executing: `python tsne_torch.py`
#
#
#  Created by <NAME> on 23-03-2020.
#  Copyright (c) 2020. All rights reserved.
import numpy as np
import matplotlib.pyplot as pyplot
import argparse
import torch
import copy
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser()
parser.add_argument("--feature", type=str, default="select_features.pt", help="torch feature file name")
parser.add_argument("--label", type=str, default="select_labels.npy", help="numpy label file name")
parser.add_argument("--cuda", type=int, default=1, help="if use cuda accelarate")
#parser.add_argument("--gpu_id", type=int, default=0, help="gpu id")
parser.add_argument("--max_iter", type=int, default=1000, help="Maximum number of iterations for the optimization")
parser.add_argument("--save_fig_name", type=str, default='result.png', help="fig name")

opt = parser.parse_args()
print("get choice from args", opt)
feature = opt.feature
label = opt.label
max_iter = opt.max_iter
cuda = opt.cuda
save_fig_name = opt.save_fig_name
#gpu_id = opt.gpu_id

if cuda:
	print("set use cuda")
	torch.set_default_tensor_type(torch.cuda.DoubleTensor)
else:
	torch.set_default_tensor_type(torch.DoubleTensor)

def Hbeta_torch(D, beta=1.0):
	P = torch.exp(-D.clone() * beta)

	sumP = torch.sum(P)

	H = torch.log(sumP) + beta * torch.sum(D * P) / sumP
	P = P / sumP

	return H, P
 
def x2p_torch(X, tol=1e-5, perplexity=30.0):
	"""
		Performs a binary search to get P-values in such a way that each
		conditional Gaussian has the same perplexity.
	"""

	# Initialize some variables
	print("Computing pairwise distances...")
	(n, d) = X.shape

	sum_X = torch.sum(X*X, 1)
	D = torch.add(torch.add(-2 * torch.mm(X, X.t()), sum_X).t(), sum_X)

	P = torch.zeros(n, n)
	beta = torch.ones(n, 1)
	logU = torch.log(torch.tensor([perplexity]))
	n_list = [i for i in range(n)]

	# Loop over all datapoints
	for i in range(n):

		# Print progress
		if i % 500 == 0:
			print("Computing P-values for point %d of %d..." % (i, n))

		# Compute the Gaussian kernel and entropy for the current precision
		# there may be something wrong with this setting None
		betamin = None
		betamax = None
		Di = D[i, n_list[0:i]+n_list[i+1:n]]

		(H, thisP) = Hbeta_torch(Di, beta[i])

		# Evaluate whether the perplexity is within tolerance
		Hdiff = H - logU
		tries = 0
		while torch.abs(Hdiff) > tol and tries < 50:

			# If not, increase or decrease precision
			if Hdiff > 0:
				betamin = beta[i].clone()
				if betamax is None:
					beta[i] = beta[i] * 2.
				else:
					beta[i] = (beta[i] + betamax) / 2.
			else:
				betamax = beta[i].clone()
				if betamin is None:
					beta[i] = beta[i] / 2.
				else:
					beta[i] = (beta[i] + betamin) / 2.

			# Recompute the values
			(H, thisP) = Hbeta_torch(Di, beta[i])

			Hdiff = H - logU
			tries += 1

		# Set the final row of P
		P[i, n_list[0:i]+n_list[i+1:n]] = thisP

	# Return final P-matrix
	return P


def pca_torch(X, no_dims=50):
	print("Preprocessing the data using PCA...")
	(n, d) = X.shape
	X = X - torch.mean(X, 0)

	(l, M) = torch.eig(torch.mm(X.t(), X), True)
	# split M real
	for i in range(d):
		if l[i, 1] != 0:
			M[:, i+1] = M[:, i]
			i += 1

	Y = torch.mm(X, M[:, 0:no_dims])
	return Y


def tsne(X, max_iter, no_dims=2, initial_dims=50, perplexity=30.0):
	"""
		Runs t-SNE on the dataset in the NxD array X to reduce its
		dimensionality to no_dims dimensions. The syntaxis of the function is
		`Y = tsne.tsne(X, no_dims, perplexity), where X is an NxD NumPy array.
	"""

	# Check inputs
	if isinstance(no_dims, float):
		print("Error: array X should not have type float.")
		return -1
	if round(no_dims) != no_dims:
		print("Error: number of dimensions should be an integer.")
		return -1

	# Initialize variables
	X = pca_torch(X, initial_dims)
	(n, d) = X.shape
	initial_momentum = 0.5
	final_momentum = 0.8
	eta = 500
	min_gain = 0.01
	Y = torch.randn(n, no_dims)
	dY = torch.zeros(n, no_dims)
	iY = torch.zeros(n, no_dims)
	gains = torch.ones(n, no_dims)

	# Compute P-values
	P = x2p_torch(X, 1e-5, perplexity)
	P = P + P.t()
	P = P / torch.sum(P)
	P = P * 4.    # early exaggeration
	print("get P shape", P.shape)
	P = torch.max(P, torch.tensor([1e-21]))

	# Run iterations
	for iter in range(max_iter):

		# Compute pairwise affinities
		sum_Y = torch.sum(Y*Y, 1)
		num = -2. * torch.mm(Y, Y.t())
		num = 1. / (1. + torch.add(torch.add(num, sum_Y).t(), sum_Y))
		num[range(n), range(n)] = 0.
		Q = num / torch.sum(num)
		Q = torch.max(Q, torch.tensor([1e-12]))

		# Compute gradient
		PQ = P - Q
		for i in range(n):
			dY[i, :] = torch.sum((PQ[:, i] * num[:, i]).repeat(no_dims, 1).t() * (Y[i, :] - Y), 0)

		# Perform the update
		if iter < 20:
			momentum = initial_momentum
		else:
			momentum = final_momentum

		gains = (gains + 0.2) * ((dY > 0.) != (iY > 0.)).double() + (gains * 0.8) * ((dY > 0.) == (iY > 0.)).double()
		gains[gains < min_gain] = min_gain
		iY = momentum * iY - eta * (gains * dY)
		Y = Y + iY
		Y = Y - torch.mean(Y, 0)

		# Compute current value of cost function
		if (iter + 1) % 10 == 0:
			C = torch.sum(P * torch.log(P / Q))
			print("Iteration %d: error is %f" % (iter + 1, C))

		# Stop lying about P-values
		if iter == 100:
			P = P / 4.

	# Return solution
	return Y


if __name__ == "__main__":
	print("Run Y = tsne.tsne(X, no_dims, perplexity) to perform t-SNE on your dataset.")

	color_set = ['red', 'green', 'blue', 'yellow', 'black', 'white']	
	data_path = './data/'
	
	X = torch.load(data_path + feature)
	if cuda: X = X.to('cuda')
	print('X size:', X.size())
	labels = np.load(data_path + label)
	print('labels size:', labels.shape)
	
	#if str(type(labels[0])) != "<class 'int'>":
	category_set = np.unique(labels).tolist()
	print('category_set:', category_set)
	
	category_index_list = []
	for i in range(len(category_set)):
		category_index_list.append(np.where(labels==category_set[i]))

	# confirm that x file get same number point than label file
	# otherwise may cause error in scatter
	assert(X.size()[0]==labels.shape[0])

	with torch.no_grad():
		Y = tsne(X, max_iter, 2, 50, 20.0)

	if opt.cuda:
		Y = Y.cpu().numpy()
	
	Y_new_list = []
	label_new_list = []
	for category_index in category_index_list:
		category_index = category_index[0].tolist()
		print(category_index)
		Y_new = []
		label_new = []
		for i in category_index:
			Y_new.append(Y[i])
			label_new.append(color_set[category_set.index(labels[i])])
		Y_new_list.append(np.array(Y_new))
		label_new_list.append(np.array(label_new))
		
	plt.figure()
	for i in range(len(category_set)):
		plt.scatter(Y_new_list[i][:, 0], Y_new_list[i][:, 1], c=label_new_list[i], label=category_set[i])
	plt.legend(fontsize='x-large')
 
	plt.savefig('./' + save_fig_name)
	#plt.show()