"""
Author: Pinar Demetci
Principal Investigator: Ritambhara Singh, Ph.D. from Brown University
08 August 2021
Updated: 23 February 2023
SCOTv2 algorithm: Single Cell alignment using Optimal Transport version 2
Correspondence: pinar_demetci@brown.edu, ritambhara@brown.edu
"""

### Import python packages we depend on:
import numpy as np
import torch
import ot
import scipy
# 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 SCOTv2(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, data):

		assert type(data)==list and len(data)>=2, "As input, SCOTv2 requires a list, containing at least two numpy arrays to be aligned.  \
				Each numpy array/matrix corresponds to a dataset, with samples (cells) in rows and features (latent representations or genomic features) in columns. \
				We recommend using latent representations (e.g. principal components for RNA-seq and topics - via cisTopic- for ATAC-seq/Methyl-seq)."
		self.data=data
		self.marginals=[] # Holds the empirical probability distributions over samples in each dataset
		self.graphs=[] # Holds graphs per dataset
		self.graphDists=[] # Holds intra-domain graph distances for each input dataset
		self.couplings=[] # Holds coupling matrices
		self.gwdists=[] # Gromov-Wasserstein distances between domains after alignment
		self.flags = [] # Holds alignment convergence flags (booleans: True/False)

		self.aligned_data=[]

	def _init_marginals(self):
		# Without any prior information, we set the probabilities to what we observe empirically: uniform over all observed sample
		for i in range(len(self.data)):
			num_cells=self.data[i].shape[0]
			marginalDist=torch.ones(num_cells)/num_cells
			self.marginals.append(marginalDist)
		return self.marginals

	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 normalized data and set the argument 'normalize=False' when running the algorithm. \
		 We have found l2 normalization to empirically perform better with single-cell sequencing datasets, including when using latent representations. "

		for i in range(len(self.data)):
			if norm=="zscore":
				scaler=StandardScaler()
				self.data[i]=scaler.fit_transform(self.data[i])
			else:
				if (bySample==True or bySample==None):
					axis=1
				else:
					axis=0
				self.data[i] =normalize(self.data[i], norm=norm, axis=axis)
		return self.data # Normalized data

	def construct_graph(self, k=20, 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

		for i in range(len(self.data)):
			self.graphs.append(kneighbors_graph(self.data[i], n_neighbors=k, mode=mode, metric=metric, include_self=include_self))

		return self.graphs

	def init_graph_distances(self):
		for i in range(len(self.data)):
			# Compute shortest distances
			shortestPath=dijkstra(csgraph= csr_matrix(self.graphs[i]), directed=False, return_predecessors=False)
			# Deal with unconnected stuff (infinities):
			Max_dist=np.nanmax(shortestPath[shortestPath != np.inf])
			shortestPath[shortestPath > Max_dist] = Max_dist
			# Finnally, normalize the distance matrix:
			self.graphDists.append(shortestPath/shortestPath.max())

		return self.graphDists

	def _exp_sinkhorn_solver(self, ecost, u, v,a,b, mass, eps, rho, rho2, nits_sinkhorn, tol_sinkhorn):
			"""
			Parameters
			----------
			- ecost: torch.Tensor of size [size_X, size_Y]
					 Exponential kernel generated from the local cost based on the current coupling.  
			- u: torch.Tensor of size [size_X[0]].
				 First dual potential defined on X.
			- v: torch.Tensor of size [size_Y[0]].
				 Second dual potential defined on Y. 
			- mass: torch.Tensor of size [1]. 
					Mass of the current coupling.
			- nits_sinkhorn: int. 
							 Maximum number of iterations to update Sinkhorn potentials in inner loop.
			- tol_sinkhorn: float
							Tolerance on convergence of Sinkhorn potentials.

			Returns
			----------
			u: torch.Tensor of size [size_X[0]]
			   First dual potential of Sinkhorn algorithm
			v: torch.Tensor of size [size_Y[0]]
			   Second dual potential of Sinkhorn algorithm
			logpi: torch.Tensor of size [size_X, size_Y]
				   Optimal transport plan in log-space.
			"""
			# Initialize potentials by finding best translation
			if u is None or v is None:
				u, v = torch.ones_like(a), torch.ones_like(b)
			k = (a * u ** (-eps / rho)).sum()+ (b * v ** (-eps / rho)).sum()
			k = k / (2 * (u[:, None] * v[None, :] * ecost *a[:, None] * b[None, :]).sum())
			z = (0.5 * mass * eps) / (2.0 + 0.5 * (eps / rho) + 0.5 * (eps / rho2))
			k = k ** z
			u,v= u * k, v * k

			# perform Sinkhorn updates in LSE form
			for j in range(nits_sinkhorn):
				u_prev = u.clone()
				v = torch.einsum("ij,i->j", ecost, a * u) ** (-1.0 / (1.0 + eps / rho))
				u = torch.einsum("ij,j->i", ecost, b * v) ** (-1.0 / (1.0 + eps / rho2))
				if (u.log() - u_prev.log()).abs().max().item() * eps < tol_sinkhorn:
					break
			pi = u[:, None] * v[None, :] * ecost * a[:, None] * b[None, :]
			return u, v, pi

	def exp_unbalanced_gw(self,a, dx, b, dy, eps=0.01, rho=1.0, rho2=None, nits_plan=3000, tol_plan=1e-6, nits_sinkhorn=3000, tol_sinkhorn=1e-6):
		if rho2 is None:
			rho2 = rho #KL divergence coefficient doesn't have to be the same for both couplings. 
					   #But, to keep #hyperparameters low, we default to using the same coefficient. 
					   #Someone else playing with our code could assign a rho2 different than rho, though.

		# Initialize the coupling and local costs
		pi= a[:, None]* b[None, :] / (a.sum() * b.sum()).sqrt()
		pi_prev = torch.zeros_like(pi)
		up, vp = None, None

		for i in range(nits_plan):
			pi_prev = pi.clone()
			mp = pi.sum()

			#Compute the current local cost:
			distxy = torch.einsum("ij,kj->ik", dx, torch.einsum("kl,jl->kj", dy, pi))
			kl_pi = torch.sum(pi * (pi / (a[:, None] * b[None, :]) + 1e-10).log())
			mu, nu = torch.sum(pi, dim=1), torch.sum(pi, dim=0)
			distxx = torch.einsum("ij,j->i", dx ** 2, mu)
			distyy = torch.einsum("kl,l->k", dy ** 2, nu)
			lcost = (distxx[:, None] + distyy[None, :] - 2 * distxy) + eps * kl_pi
			if rho < float("Inf"):
				lcost = (lcost+ rho* torch.sum(mu * (mu / a + 1e-10).log()))
			if rho2 < float("Inf"):
				lcost = (lcost+ rho2* torch.sum(nu * (nu / b + 1e-10).log()))
			ecost = (-lcost /(mp * eps)).exp()

			if (i%10)==0:
				print("Unbalanced GW step:", i)
			#compute the coupling via sinkhorn
			up, vp, pi = self._exp_sinkhorn_solver(ecost, up, vp, a, b, mp, eps, rho, rho2,nits_sinkhorn, tol_sinkhorn)
			
			flag=True
			if torch.any(torch.isnan(pi)):
				flag=False

			pi = (mp / pi.sum()).sqrt() * pi
			if (pi - pi_prev).abs().max().item() < tol_plan:
				break
		return pi, flag

	def find_correspondences(self, normalize=True, norm="l2", bySample=True, k=20, mode= "connectivity", metric="correlation",  eps=0.01, rho=1.0, rho2=None):
		# Normalize 
		if normalize:
			self._normalize(norm=norm, bySample=bySample)
		# Initialize inputs for (unbalanced) Gromov-Wasserstein optimal transport:
		self._init_marginals()
		print("computing intra-domain graph distances")
		self.construct_graph(k=k, mode=mode, metric=metric)
		self.init_graph_distances()
		# Run pairwise dataset alignments:
		for i in range(len(self.data)-1):
			print("running pairwise dataset alignments")
			a,b =torch.Tensor(self.marginals[0]), torch.Tensor(self.marginals[i+1])
			dx, dy= torch.Tensor(self.graphDists[0]), torch.Tensor(self.graphDists[i+1])
			coupling, flag=self.exp_unbalanced_gw(a, dx, b, dy, eps=eps, rho=rho, rho2=rho2, nits_plan=3000, tol_plan=1e-6, nits_sinkhorn=3000, tol_sinkhorn=1e-6)
			self.couplings.append(coupling)
			self.flags.append(flag)
			if flag==False:
					raise Exception(
					f"Solver got NaN plan with params (eps, rho, rho2) "
					f" = {eps, rho, rho2}. Try increasing argument eps")
		return self.couplings

	def barycentric_projection(self):
		aligned_datasets=[self.data[0]]
		for i in range(0,len(self.couplings)):
			coupling=np.transpose(self.couplings[i].numpy())
			weights=np.sum(coupling, axis = 1)
			projected_data=np.matmul((coupling/ weights[:, None]), self.data[0])
			aligned_datasets.append(projected_data)
		return aligned_datasets

	def coembed_datasets(self, Lambda=1.0, out_dim=10):
		"""
		Co-embeds datasets in a shared space.
		Implementation is based on Cao et al 2022 (Pamona)
		"""
		n_datasets = len(self.data)
		H0 = []
		L = []
		for i in range(n_datasets-1):
			self.couplings[i] = self.couplings[i]*np.shape(self.data[i])[0]

		for i in range(n_datasets):    
			graph_data = self.graphs[i] + self.graphs[i].T.multiply(self.graphs[i].T > self.graphs[i]) - \
				self.graphs[i].multiply(self.graphs[i].T > self.graphs[i])
			W = np.array(graph_data.todense())
			index_pos = np.where(W>0)
			W[index_pos] = 1/W[index_pos] 
			D = np.diag(np.dot(W, np.ones(np.shape(W)[1])))
			L.append(D - W)

		Sigma_x = []
		Sigma_y = []
		for i in range(n_datasets-1):
			Sigma_y.append(np.diag(np.dot(np.transpose(np.ones(np.shape(self.couplings[i])[0])), self.couplings[i])))
			Sigma_x.append(np.diag(np.dot(self.couplings[i], np.ones(np.shape(self.couplings[i])[1]))))

		S_xy = self.couplings[0]
		S_xx = L[0] + Lambda*Sigma_x[0]
		S_yy = L[-1] +Lambda*Sigma_y[0]
		for i in range(1, n_datasets-1):
			S_xy = np.vstack((S_xy, self.couplings[i]))
			S_xx = block_diag(S_xx, L[i] + Lambda*Sigma_x[i])
			S_yy = S_yy + Lambda*Sigma_y[i]

		v, Q = np.linalg.eig(S_xx)
		v = v + 1e-12   
		V = np.diag(v**(-0.5))
		H_x = np.dot(Q, np.dot(V, np.transpose(Q)))

		v, Q = np.linalg.eig(S_yy)
		v = v + 1e-12      
		V = np.diag(v**(-0.5))
		H_y = np.dot(Q, np.dot(V, np.transpose(Q)))

		H = np.dot(H_x, np.dot(S_xy, H_y))
		U, sigma, V = np.linalg.svd(H)

		num = [0]
		for i in range(n_datasets-1):
			num.append(num[i]+len(self.data[i]))

		U, V = U[:,:out_dim], np.transpose(V)[:,:out_dim]

		fx = np.dot(H_x, U)
		fy = np.dot(H_y, V)

		integrated_data = []
		for i in range(n_datasets-1):
			integrated_data.append(fx[num[i]:num[i+1]])

		integrated_data.append(fy)

		return integrated_data

	def align(self,normalize=True, norm="l2", bySample=True, k=20, mode= "connectivity", metric="correlation",  eps=0.01, rho=1.0, rho2=None, projMethod="embedding", Lambda=1.0, out_dim=10):
		assert projMethod in ["embedding", "barycentric"], "The input to the parameter 'projMethod' needs to be one of \
								'embedding' (if co-embedding them in a new shared space) or 'barycentric' (if using barycentric projection)"
		self.find_correspondences(normalize=normalize, norm=norm, bySample=bySample, k=k, mode=mode, metric=metric,  eps=eps, rho=rho, rho2=rho2)
		print("FLAGS", self.flags)
		if projMethod=="embedding":
			integrated_data=self.coembed_datasets(Lambda=Lambda, out_dim=out_dim)
		else:
			integrated_data=self.barycentric_projection()
		self.integrated_data=integrated_data
		return integrated_data
	
# X=np.load("../data/SNARE/SNAREseq_atac_feat.npy")[0:1000,:]
# Y=np.load("../data/SNARE/SNAREseq_rna_feat.npy")
# print(X.shape, Y.shape)
# SCOT=SCOTv2([Y,X])
# aligned_datasets=SCOT.align(normalize=True, k=50, eps=0.005, rho=0.1, projMethod="barycentric")
# print(len(aligned_datasets))
# print(aligned_datasets[0].shape)
# print(aligned_datasets[1].shape)
# # np.save("aligned_atac.npy", aligned_datasets[1])
# np.save("aligned_rna.npy", aligned_datasets[0])