"""
Authors: Pinar Demetci, Rebecca Santorella
Principal Investigator: Ritambhara Singh, Ph.D. from Brown University
12 February 2020
Updated: 27 November 2020
SCOT algorithm (version 1): Single Cell alignment using Optimal Transport
Correspondence: pinar_demetci@brown.edu, rebecca_santorella@brown.edu, ritambhara@brown.edu
"""

### Import python packages we depend on:
# For regular matrix operations:
import numpy as np
# For optimal transport operations:
import ot
# For computing graph distances:
from scipy.sparse.csgraph import dijkstra
from scipy.sparse import csr_matrix
from sklearn.neighbors import kneighbors_graph

# For pre-processing, normalization
from sklearn.preprocessing import StandardScaler, normalize


class SCOT(object):
	"""
	SCOT algorithm for unsupervised alignment of single-cell multi-omic data.
	https://www.biorxiv.org/content/10.1101/2020.04.28.066787v2 (original preprint)
	https://www.liebertpub.com/doi/full/10.1089/cmb.2021.0446 (Journal of Computational Biology publication through RECOMB 2021 conference)
	Input: domain1, domain2 in form of numpy arrays/matrices, where the rows correspond to samples and columns correspond to features.
	Returns: aligned domain 1, aligned domain 2 in form of numpy arrays/matrices projected on domain 1
	Example use:
	# Given two numpy matrices, domain1 and domain2, where the rows are cells and columns are different genomic features:
	scot= SCOT(domain1, domain2)
	aligned_domain1, aligned_domain2 = scot.align(k=20, e=1e-3)
	#If you can't pick the parameters k and e, you can try out our unsupervised self-tuning heuristic by running:
	scot= SCOT(domain1, domain2)
	aligned_domain1, aligned_domain2 = scot.align(selfTune=True)
	Required parameters:
	- k: Number of neighbors to be used when constructing kNN graphs. Default= min(min(n_1, n_2), 50), where n_i, for i=1,2 corresponds to the number of samples in the i^th domain.
	- e: Regularization constant for the entropic regularization term in entropic Gromov-Wasserstein optimal transport formulation. Default= 1e-3 
   
	Optional parameters:
	- normalize= Determines whether to normalize input data ahead of alignment. True or False (boolean parameter). Default = True.
	- norm= Determines what sort of normalization to run, "l2", "l1", "max", "zscore". Default="l2" 
	- mode: "connectivity" or "distance". Determines whether to use a connectivity graph (adjacency matrix of 1s/0s based on whether nodes are connected) or a distance graph (adjacency matrix entries weighted by distances between nodes). Default="connectivity"  
	- metric: Sets the metric to use while constructing nearest neighbor graphs. some possible choices are "correlation", "minkowski".  "correlation" is Pearson's correlation and "minkowski" is equivalent to Euclidean distance in its default form (). Default= "correlation". 
	- verbose: Prints loss while optimizing the optimal transport formulation. Default=True
	- XontoY: Determines the direction of barycentric projection. True or False (boolean parameter). If True, projects domain1 onto domain2. If False, projects domain2 onto domain1. Default=True.
	Note: If you want to specify the marginal distributions of the input domains and not use uniform distribution, please set the attributes p and q to the distributions of your choice (for domain 1, and 2, respectively) 
			after initializing a SCOT class instance and before running alignment and set init_marginals=False in .align() parameters
	"""

	def __init__(self, domain1, domain2):

		self.X=domain1
		self.y=domain2

		self.p= None #empirical probability distribution for domain 1 (X)
		self.q= None #empirical probability distribution for domain 2 (y)

		self.Cx=None #intra-domain graph distances for domain 1 (X)
		self.Cy=None #intra-domain graph distances for domain 2 (y)

		self.coupling=None # Coupling matrix that relates domain 1 and domain 2, ..., m
		self.gwdist=None # Gromov-Wasserstein distance between domains after alignment
		self.flag = None # convergence flag

		self.X_aligned=None #aligned datasets to return: domain1
		self.y_aligned=None #aligned datasets to return: domain2

	def init_marginals(self):
		# Without any prior information, we set the probabilities to what we observe empirically: uniform over all observed sample
		self.p= ot.unif(self.X.shape[0])
		self.q = ot.unif(self.y.shape[0])

	def normalize(self, norm="l2", bySample=True):
		assert (norm in ["l1","l2","max", "zscore"]), "Norm argument has to be either one of 'max', 'l1', 'l2' or 'zscore'. If you would like to perform another type of normalization, please give SCOT the normalize data and set the argument normalize=False when running the algorithm."

		if (bySample==True or bySample==None):
			axis=1
		else:
			axis=0

		if norm=="zscore":
			scaler=StandardScaler()
			self.X, self.y=scaler.fit_transform(self.X), scaler.fit_transform(self.y)

		else:
			self.X, self.y =normalize(self.X, norm=norm, axis=axis), normalize(self.y, norm=norm, axis=axis)

	def construct_graph(self, k, mode= "connectivity", metric="correlation"):
		assert (mode in ["connectivity", "distance"]), "Norm argument has to be either one of 'connectivity', or 'distance'. "
		if mode=="connectivity":
			include_self=True
		else:
			include_self=False

		self.Xgraph=kneighbors_graph(self.X, k, mode=mode, metric=metric, include_self=include_self)
		self.ygraph=kneighbors_graph(self.y, k, mode=mode, metric=metric, include_self=include_self)

		return self.Xgraph, self.ygraph

	def init_distances(self):
		# Compute shortest distances
		X_shortestPath=dijkstra(csgraph= csr_matrix(self.Xgraph), directed=False, return_predecessors=False)
		y_shortestPath=dijkstra(csgraph= csr_matrix(self.ygraph), directed=False, return_predecessors=False)

		# Deal with unconnected stuff (infinities):
		X_max=np.nanmax(X_shortestPath[X_shortestPath != np.inf])
		y_max=np.nanmax(y_shortestPath[y_shortestPath != np.inf])
		X_shortestPath[X_shortestPath > X_max] = X_max
		y_shortestPath[y_shortestPath > y_max] = y_max

		# Finnally, normalize the distance matrix:
		self.Cx=X_shortestPath/X_shortestPath.max()
		self.Cy=y_shortestPath/y_shortestPath.max()

		return self.Cx, self.Cy

	def find_correspondences(self, e, verbose=True):
		self.coupling, log= ot.gromov.entropic_gromov_wasserstein(self.Cx, self.Cy, self.p, self.q, loss_fun='square_loss', epsilon=e, log=True, verbose=verbose)
		self.gwdist=log['gw_dist']

		# Check convergence:
		if (np.isnan(self.coupling).any() or np.any(~self.coupling.any(axis=1)) or np.any(~self.coupling.any(axis=0)) or sum(sum(self.coupling)) < .95):
			self.flag=False
		else:
			self.flag=True

		return self.gwdist

	def barycentric_projection(self, XontoY=True):
		if XontoY:
			#Projecting the first domain onto the second domain
			self.y_aligned=self.y
			weights=np.sum(self.coupling, axis = 1)
			self.X_aligned=np.matmul(self.coupling, self.y) / weights[:, None]
		else:
			#Projecting the second domain onto the first domain
			self.X_aligned=self.X
			weights=np.sum(self.coupling, axis = 0)
			self.y_aligned=np.matmul(np.transpose(self.coupling), self.X) / weights[:, None]
		return self.X_aligned, self.y_aligned

	def align(self, k=None, e=1e-3, mode="connectivity", metric="correlation", verbose=True, normalize=True, norm="l2", XontoY=True, selfTune=False, init_marginals=True):
		if normalize:
			self.normalize(norm=norm)
		if init_marginals:
			self.init_marginals()

		if selfTune:
			X_aligned, y_aligned= self.unsupervised_scot()
		else:
			if k==None:
				k=min((int(self.X.shape[0]*0.2), int(self.y.shape[0]*0.2)),50)

			self.construct_graph(k, mode= "connectivity", metric="correlation")
			self.init_distances()
			self.find_correspondences(e=e, verbose=verbose)

			if self.flag==False:
				print("CONVERGENCE ERROR: Optimization procedure runs into numerical errors with the hyperparameters specified. Please try aligning with higher values of epsilon.")
				return
			
			else:
				X_aligned, y_aligned = self.barycentric_projection(XontoY=XontoY)

		self.X_aligned, self.y_aligned=X_aligned, y_aligned
		return self.X_aligned, self.y_aligned

	def search_scot(self, ks, es, all_values = False,  mode= "connectivity", metric="correlation", normalize=True, norm="l2", init_marginals=True):
		'''
		Performs a hyperparameter sweep for given values of k and epsilon
		Default: return the parameters corresponding to the lowest GW distance
		(Optional): return all k, epsilon, and GW values
		'''

		# initialize alignment
		if normalize:
			self.normalize(norm=norm)
		if init_marginals:
			self.init_marginals()

		# Note to self: Incorporate multiprocessing here to speed things up
		# store values of k, epsilon, and gw distance
		total=len(es)*len(ks)
		k_sweep=np.zeros(total)
		e_sweep=np.zeros(total)
		gw_sweep=np.zeros(total)

		gmin = 1
		counter=1

		X_aligned,y_aligned=None, None
		e_best,k_best=None, None
		# search in k first to reduce graph computation
		for k in ks:
			self.construct_graph(k, mode= mode, metric=metric)
			self.init_distances()
			for e in es:
				print(counter, "/", total)
				print("Aligning k: ",k, " and e: ",e)
				# run alignment / optimize correspondence matrix:
				self.find_correspondences(e=e, verbose=False)
				# save values
				if self.flag:
					if all_values:
						k_sweep[counter]=k
						e_sweep[counter]=e
						gw_sweep[counter] = self.gwdist

						print(self.gwdist)
					# save the alignment if it is lower
					if self.gwdist < gmin:
						X_aligned, y_aligned = self.barycentric_projection()
						gmin =self.gwdist
						e_best, k_best= e, k
					counter = counter + 1
		   
		if all_values:
			# return alignment and all values
			return X_aligned, y_aligned, gw_sweep, k_sweep, e_sweep
		else:
			# return  alignment and the parameters corresponding to the lowest GW distance
			return X_aligned, y_aligned, gmin, k_best, e_best


	def unsupervised_scot(self, normalize=False, norm='l2'):
		'''
		Unsupervised hyperparameter tuning algorithm to find an alignment
		by using the GW distance as a measure of alignment
		'''

		# use k = 20% of # sample or k = 50 if dataset is large
		n = min(self.X.shape[0], self.y.shape[0])
		k_start = min(n // 5, 50)

		num_eps = 12
		num_k = 5

		# define search space
		es = np.logspace(-1, -3, num_eps)
		if ( n > 250):
			ks = np.linspace(20, 100, num_k)
		else:
			ks = np.linspace(n//20, n//6, num_k)
		ks = ks.astype(int)
		
		# search parameter space
		X_aligned, y_aligned, g_best, k_best, e_best = self.search_scot(ks, es, all_values=False, normalize=normalize, norm=norm, init_marginals=False)

		print("Alignment completed. Hyperparameters selected from the unsupervised hyperparameter sweep are: %d for number of neighbors k and %f for epsilon" %(k_best, e_best))

		return X_aligned,