from typing import Optional, Union

import warnings

import os

import numpy as np

import pandas as pd

from scipy import stats

import VITAE.model as model 

import VITAE.train as train 

from VITAE.inference import Inferer

from VITAE.utils import get_igraph, leidenalg_igraph, \

   DE_test, _comp_dist, _get_smooth_curve

from VITAE.metric import topology, get_GRI

import tensorflow as tf

from sklearn.metrics.cluster import adjusted_rand_score

from sklearn.model_selection import train_test_split

from sklearn.cluster import AgglomerativeClustering

from sklearn.preprocessing import OneHotEncoder, StandardScaler, OrdinalEncoder

import scanpy as sc

import networkx as nx

import matplotlib.pyplot as plt

import matplotlib.patheffects as pe

class VITAE():

    """

    Variational Inference for Trajectory by AutoEncoder.

    """

    def __init__(self, adata: sc.AnnData,

               covariates = None, pi_covariates = None,

               model_type: str = 'Gaussian',

               npc: int = 64,

               adata_layer_counts = None,

               copy_adata: bool = False,

               hidden_layers = [32],

               latent_space_dim: int = 16,

               conditions = None):

        '''

        Get input data for model. Data need to be first processed using scancy and stored as an AnnData object

         The 'UMI' or 'non-UMI' model need the original count matrix, so the count matrix need to be saved in

         adata.layers in order to use these models.

        Parameters

        ----------

        adata : sc.AnnData

            The scanpy AnnData object. adata should already contain adata.var.highly_variable

        covariates : list, optional

            A list of names of covariate vectors that are stored in adata.obs

        pi_covariates: list, optional

            A list of names of covariate vectors used as input for pilayer

        model_type : str, optional

            'UMI', 'non-UMI' and 'Gaussian', default is 'Gaussian'.

        npc : int, optional

            The number of PCs to use when model_type is 'Gaussian'. The default is 64.

        adata_layer_counts: str, optional

            the key name of adata.layers that stores the count data if model_type is

            'UMI' or 'non-UMI'

        copy_adata: bool, optional. Set to True if we don't want VITAE to modify the original adata. If set to True, self.adata will be an independent copy of the original adata. 

        hidden_layers : list, optional

            The list of dimensions of layers of autoencoder between latent space and original space. Default is to have only one hidden layer with 32 nodes

        latent_space_dim : int, optional

            The dimension of latent space.

        gamme : float, optional

            The weight of the MMD loss

        conditions : str or list, optional

            The conditions of different cells

        Returns

        -------

        None.

        '''

        self.dict_method_scname = {

            'PCA' : 'X_pca',

            'UMAP' : 'X_umap',

            'TSNE' : 'X_tsne',

            'diffmap' : 'X_diffmap',

            'draw_graph' : 'X_draw_graph_fa'

        }

        if model_type != 'Gaussian':

            if adata_layer_counts is None:

                raise ValueError("need to provide the name in adata.layers that stores the raw count data")

            if 'highly_variable' not in adata.var:

                raise ValueError("need to first select highly variable genes using scanpy")

        self.model_type = model_type

        if copy_adata:

            self.adata = adata.copy()

        else:

            self.adata = adata

        if covariates is not None:

            if isinstance(covariates, str):

                covariates = [covariates]

            covariates = np.array(covariates)

            id_cat = (adata.obs[covariates].dtypes == 'category')

            # add OneHotEncoder & StandardScaler as class variable if needed

            if np.sum(id_cat)>0:

                covariates_cat = OneHotEncoder(drop='if_binary', handle_unknown='ignore'

                    ).fit_transform(adata.obs[covariates[id_cat]]).toarray()

            else:

                covariates_cat = np.array([]).reshape(adata.shape[0],0)

            # temporarily disable StandardScaler

            if np.sum(~id_cat)>0:

                #covariates_con = StandardScaler().fit_transform(adata.obs[covariates[~id_cat]])

                covariates_con = adata.obs[covariates[~id_cat]]

            else:

                covariates_con = np.array([]).reshape(adata.shape[0],0)

            self.covariates = np.c_[covariates_cat, covariates_con].astype(tf.keras.backend.floatx())

        else:

            self.covariates = None

        if conditions is not None:

            ## observations with np.nan will not participant in calculating mmd_loss

            if isinstance(conditions, str):

                conditions = [conditions]

            conditions = np.array(conditions)

            if np.any(adata.obs[conditions].dtypes != 'category'):

                raise ValueError("Conditions should all be categorical.")

            self.conditions = OrdinalEncoder(dtype=int, encoded_missing_value=-1).fit_transform(adata.obs[conditions]) + int(1)

        else:

            self.conditions = None

        if pi_covariates is not None:

            self.pi_cov = adata.obs[pi_covariates].to_numpy()

            if self.pi_cov.ndim == 1:

                self.pi_cov = self.pi_cov.reshape(-1, 1)

                self.pi_cov = self.pi_cov.astype(tf.keras.backend.floatx())

        else:

            self.pi_cov = np.zeros((adata.shape[0],1), dtype=tf.keras.backend.floatx())

        self.model_type = model_type

        self._adata = sc.AnnData(X = self.adata.X, var = self.adata.var)

        self._adata.obs = self.adata.obs

        self._adata.uns = self.adata.uns

        if model_type == 'Gaussian':

            sc.tl.pca(adata, n_comps = npc)

            self.X_input = self.X_output = adata.obsm['X_pca']

            self.scale_factor = np.ones(self.X_output.shape[0])

        else:

            print(f"{adata.var.highly_variable.sum()} highly variable genes selected as input") 

            self.X_input = adata.X[:, adata.var.highly_variable]

            self.X_output = adata.layers[adata_layer_counts][ :, adata.var.highly_variable]

            self.scale_factor = np.sum(self.X_output, axis=1, keepdims=True)/1e4

        self.dimensions = hidden_layers

        self.dim_latent = latent_space_dim

        self.vae = model.VariationalAutoEncoder(

            self.X_output.shape[1], self.dimensions,

            self.dim_latent, self.model_type,

            False if self.covariates is None else True,

            )

        if hasattr(self, 'inferer'):

            delattr(self, 'inferer')

    def pre_train(self, test_size = 0.1, random_state: int = 0,

            learning_rate: float = 1e-3, batch_size: int = 256, L: int = 1, alpha: float = 0.10, gamma: float = 0,

            phi : float = 1,num_epoch: int = 200, num_step_per_epoch: Optional[int] = None,

            early_stopping_patience: int = 10, early_stopping_tolerance: float = 0.01, 

            early_stopping_relative: bool = True, verbose: bool = False,path_to_weights: Optional[str] = None):

        '''Pretrain the model with specified learning rate.

        Parameters

        ----------

        test_size : float or int, optional

            The proportion or size of the test set.

        random_state : int, optional

            The random state for data splitting.

        learning_rate : float, optional

            The initial learning rate for the Adam optimizer.

        batch_size : int, optional 

            The batch size for pre-training.  Default is 256. Set to 32 if number of cells is small (less than 1000)

        L : int, optional 

            The number of MC samples.

        alpha : float, optional

            The value of alpha in [0,1] to encourage covariate adjustment. Not used if there is no covariates.

        gamma : float, optional

            The weight of the mmd loss if used.

        phi : float, optional

            The weight of Jocob norm of the encoder.

        num_epoch : int, optional 

            The maximum number of epochs.

        num_step_per_epoch : int, optional 

            The number of step per epoch, it will be inferred from number of cells and batch size if it is None.            

        early_stopping_patience : int, optional 

            The maximum number of epochs if there is no improvement.

        early_stopping_tolerance : float, optional 

            The minimum change of loss to be considered as an improvement.

        early_stopping_relative : bool, optional

            Whether monitor the relative change of loss as stopping criteria or not.

        path_to_weights : str, optional 

            The path of weight file to be saved; not saving weight if None.

        conditions : str or list, optional

            The conditions of different cells

        '''

        id_train, id_test = train_test_split(

                                np.arange(self.X_input.shape[0]), 

                                test_size=test_size, 

                                random_state=random_state)

        if num_step_per_epoch is None:

            num_step_per_epoch = len(id_train)//batch_size+1

        self.train_dataset = train.warp_dataset(self.X_input[id_train].astype(tf.keras.backend.floatx()), 

                                                None if self.covariates is None else self.covariates[id_train].astype(tf.keras.backend.floatx()),

                                                batch_size, 

                                                self.X_output[id_train].astype(tf.keras.backend.floatx()), 

                                                self.scale_factor[id_train].astype(tf.keras.backend.floatx()),

                                                conditions = None if self.conditions is None else self.conditions[id_train].astype(tf.keras.backend.floatx()))

        self.test_dataset = train.warp_dataset(self.X_input[id_test], 

                                                None if self.covariates is None else self.covariates[id_test].astype(tf.keras.backend.floatx()),

                                                batch_size, 

                                                self.X_output[id_test].astype(tf.keras.backend.floatx()), 

                                                self.scale_factor[id_test].astype(tf.keras.backend.floatx()),

                                                conditions = None if self.conditions is None else self.conditions[id_test].astype(tf.keras.backend.floatx()))

        self.vae = train.pre_train(

            self.train_dataset,

            self.test_dataset,

            self.vae,

            learning_rate,                        

            L, alpha, gamma, phi,

            num_epoch,

            num_step_per_epoch,

            early_stopping_patience,

            early_stopping_tolerance,

            early_stopping_relative,

            verbose)

        self.update_z()

        if path_to_weights is not None:

            self.save_model(path_to_weights)

    def update_z(self):

        self.z = self.get_latent_z()        

        self._adata_z = sc.AnnData(self.z)

        sc.pp.neighbors(self._adata_z)

    def get_latent_z(self):

        ''' get the posterier mean of current latent space z (encoder output)

        Returns

        ----------

        z : np.array

            \([N,d]\) The latent means.

        ''' 

        c = None if self.covariates is None else self.covariates

        return self.vae.get_z(self.X_input, c)

    def visualize_latent(self, method: str = "UMAP", 

                         color = None, **kwargs):

        '''

        visualize the current latent space z using the scanpy visualization tools

        Parameters

        ----------

        method : str, optional

            Visualization method to use. The default is "draw_graph" (the FA plot). Possible choices include "PCA", "UMAP", 

            "diffmap", "TSNE" and "draw_graph"

        color : TYPE, optional

            Keys for annotations of observations/cells or variables/genes, e.g., 'ann1' or ['ann1', 'ann2'].

            The default is None. Same as scanpy.

        **kwargs :  

            Extra key-value arguments that can be passed to scanpy plotting functions (scanpy.pl.XX).   

        Returns

        -------

        None.

        '''

        if method not in ['PCA', 'UMAP', 'TSNE', 'diffmap', 'draw_graph']:

            raise ValueError("visualization method should be one of 'PCA', 'UMAP', 'TSNE', 'diffmap' and 'draw_graph'")

        temp = list(self._adata_z.obsm.keys())

        if method == 'PCA' and not 'X_pca' in temp:

            print("Calculate PCs ...")

            sc.tl.pca(self._adata_z)

        elif method == 'UMAP' and not 'X_umap' in temp:  

            print("Calculate UMAP ...")

            sc.tl.umap(self._adata_z)

        elif method == 'TSNE' and not 'X_tsne' in temp:

            print("Calculate TSNE ...")

            sc.tl.tsne(self._adata_z)

        elif method == 'diffmap' and not 'X_diffmap' in temp:

            print("Calculate diffusion map ...")

            sc.tl.diffmap(self._adata_z)

        elif method == 'draw_graph' and not 'X_draw_graph_fa' in temp:

            print("Calculate FA ...")

            sc.tl.draw_graph(self._adata_z)

        self._adata.obs = self.adata.obs.copy()

        self._adata.obsp = self._adata_z.obsp

#        self._adata.uns = self._adata_z.uns

        self._adata.obsm = self._adata_z.obsm

        if method == 'PCA':

            axes = sc.pl.pca(self._adata, color = color, **kwargs)

        elif method == 'UMAP':            

            axes = sc.pl.umap(self._adata, color = color, **kwargs)

        elif method == 'TSNE':

            axes = sc.pl.tsne(self._adata, color = color, **kwargs)

        elif method == 'diffmap':

            axes = sc.pl.diffmap(self._adata, color = color, **kwargs)

        elif method == 'draw_graph':

            axes = sc.pl.draw_graph(self._adata, color = color, **kwargs)

        return axes

    def init_latent_space(self, cluster_label = None, log_pi = None, res: float = 1.0, 

                          ratio_prune= None, dist = None, dist_thres = 0.5, topk=0, pilayer = False):

        '''Initialize the latent space.

        Parameters

        ----------

        cluster_label : str, optional

            The name of vector of labels that can be found in self.adata.obs. 

            Default is None, which will perform leiden clustering on the pretrained z to get clusters

        mu : np.array, optional

            \([d,k]\) The value of initial \(\\mu\).

        log_pi : np.array, optional

            \([1,K]\) The value of initial \(\\log(\\pi)\).

        res: 

            The resolution of leiden clustering, which is a parameter value controlling the coarseness of the clustering. 

            Higher values lead to more clusters. Deafult is 1.

        ratio_prune : float, optional

            The ratio of edges to be removed before estimating.

        topk : int, optional

            The number of top k neighbors to keep for each cluster.

        '''   

        if cluster_label is None:

            print("Perform leiden clustering on the latent space z ...")

            g = get_igraph(self.z)

            cluster_labels = leidenalg_igraph(g, res = res)

            cluster_labels = cluster_labels.astype(str) 

            uni_cluster_labels = np.unique(cluster_labels)

        else:

            if isinstance(cluster_label,str):

                cluster_labels = self.adata.obs[cluster_label].to_numpy()

                uni_cluster_labels = np.array(self.adata.obs[cluster_label].cat.categories)

            else:

                ## if cluster_label is a list

                cluster_labels = cluster_label

                uni_cluster_labels = np.unique(cluster_labels)

        n_clusters = len(uni_cluster_labels)

        if not hasattr(self, 'z'):

            self.update_z()        

        z = self.z

        mu = np.zeros((z.shape[1], n_clusters))

        for i,l in enumerate(uni_cluster_labels):

            mu[:,i] = np.mean(z[cluster_labels==l], axis=0)

        if dist is None:

            ### update cluster centers if some cluster centers are too close

            clustering = AgglomerativeClustering(

                n_clusters=None,

                distance_threshold=dist_thres,

                linkage='complete'

                ).fit(mu.T/np.sqrt(mu.shape[0]))

            n_clusters_new = clustering.n_clusters_

            if n_clusters_new < n_clusters:

                print("Merge clusters for cluster centers that are too close ...")

                n_clusters = n_clusters_new

                for i in range(n_clusters):    

                    temp = uni_cluster_labels[clustering.labels_ == i]

                    idx = np.isin(cluster_labels, temp)

                    cluster_labels[idx] = ','.join(temp)

                    if np.sum(clustering.labels_==i)>1:

                        print('Merge %s'% ','.join(temp))

                uni_cluster_labels = np.unique(cluster_labels)

                mu = np.zeros((z.shape[1], n_clusters))

                for i,l in enumerate(uni_cluster_labels):

                    mu[:,i] = np.mean(z[cluster_labels==l], axis=0)

        self.adata.obs['vitae_init_clustering'] = cluster_labels

        self.adata.obs['vitae_init_clustering'] = self.adata.obs['vitae_init_clustering'].astype('category')

        print("Initial clustering labels saved as 'vitae_init_clustering' in self.adata.obs.")

        if (log_pi is None) and (cluster_labels is not None) and (n_clusters>3):                         

            n_states = int((n_clusters+1)*n_clusters/2)

            if dist is None:

                dist = _comp_dist(z, cluster_labels, mu.T)

            C = np.triu(np.ones(n_clusters))

            C[C>0] = np.arange(n_states)

            C = C + C.T - np.diag(np.diag(C))

            C = C.astype(int)

            log_pi = np.zeros((1,n_states))            

            ## pruning to throw away edges for far-away clusters if there are too many clusters

            if ratio_prune is not None:

                log_pi[0, C[np.triu(dist)>np.quantile(dist[np.triu_indices(n_clusters, 1)], 1-ratio_prune)]] = - np.inf

            else:

                log_pi[0, C[np.triu(dist)>np.quantile(dist[np.triu_indices(n_clusters, 1)], 5/n_clusters) * 3]] = - np.inf

            ## also keep the top k neighbor of clusters

            topk = max(0, min(topk, n_clusters-1)) + 1

            topk_indices = np.argsort(dist,axis=1)[:,:topk]

            for i in range(n_clusters):

                log_pi[0, C[i, topk_indices[i]]] = 0

        self.n_states = n_clusters

        self.labels = cluster_labels

        labels_map = pd.DataFrame.from_dict(

            {i:label for i,label in enumerate(uni_cluster_labels)}, 

            orient='index', columns=['label_names'], dtype=str

            )

        self.labels_map = labels_map

        self.vae.init_latent_space(self.n_states, mu, log_pi)

        self.inferer = Inferer(self.n_states)

        self.mu = self.vae.latent_space.mu.numpy()

        self.pi = np.triu(np.ones(self.n_states))

        self.pi[self.pi > 0] = tf.nn.softmax(self.vae.latent_space.pi).numpy()[0]

        if pilayer:

            self.vae.create_pilayer()

    def update_latent_space(self, dist_thres: float=0.5):

        pi = self.pi[np.triu_indices(self.n_states)]

        mu = self.mu    

        clustering = AgglomerativeClustering(

            n_clusters=None,

            distance_threshold=dist_thres,

            linkage='complete'

            ).fit(mu.T/np.sqrt(mu.shape[0]))

        n_clusters = clustering.n_clusters_   

        if n_clusters<self.n_states:      

            print("Merge clusters for cluster centers that are too close ...")

            mu_new = np.empty((self.dim_latent, n_clusters))

            C = np.zeros((self.n_states, self.n_states))

            C[np.triu_indices(self.n_states, 0)] = pi

            C = np.triu(C, 1) + C.T

            C_new = np.zeros((n_clusters, n_clusters))

            uni_cluster_labels = self.labels_map['label_names'].to_numpy()

            returned_order = {}

            cluster_labels = self.labels

            for i in range(n_clusters):

                temp = uni_cluster_labels[clustering.labels_ == i]

                idx = np.isin(cluster_labels, temp)

                cluster_labels[idx] = ','.join(temp)

                returned_order[i] = ','.join(temp)

                if np.sum(clustering.labels_==i)>1:

                    print('Merge %s'% ','.join(temp))

            uni_cluster_labels = np.unique(cluster_labels) 

            for i,l in enumerate(uni_cluster_labels):  ## reorder the merged clusters based on the cluster names

                k = np.where(returned_order == l)

                mu_new[:, i] = np.mean(mu[:,clustering.labels_==k], axis=-1)

                # sum of the aggregated pi's

                C_new[i, i] = np.sum(np.triu(C[clustering.labels_==k,:][:,clustering.labels_==k]))

                for j in range(i+1, n_clusters):

                    k1 = np.where(returned_order == uni_cluster_labels[j])

                    C_new[i, j] = np.sum(C[clustering.labels_== k, :][:, clustering.labels_==k1])

#            labels_map_new = {}

#            for i in range(n_clusters):                       

#                # update label map: int->str

#                labels_map_new[i] = self.labels_map.loc[clustering.labels_==i, 'label_names'].str.cat(sep=',')

#                if np.sum(clustering.labels_==i)>1:

#                    print('Merge %s'%labels_map_new[i])

#                # mean of the aggregated cluster means

#                mu_new[:, i] = np.mean(mu[:,clustering.labels_==i], axis=-1)

#                # sum of the aggregated pi's

#                C_new[i, i] = np.sum(np.triu(C[clustering.labels_==i,:][:,clustering.labels_==i]))

#                for j in range(i+1, n_clusters):

#                    C_new[i, j] = np.sum(C[clustering.labels_== i, :][:, clustering.labels_==j])

            C_new = np.triu(C_new,1) + C_new.T

            pi_new = C_new[np.triu_indices(n_clusters)]

            log_pi_new = np.log(pi_new, out=np.ones_like(pi_new)*(-np.inf), where=(pi_new!=0)).reshape((1,-1))

            self.n_states = n_clusters

            self.labels_map = pd.DataFrame.from_dict(

                {i:label for i,label in enumerate(uni_cluster_labels)},

                orient='index', columns=['label_names'], dtype=str

                )

            self.labels = cluster_labels

#            self.labels_map = pd.DataFrame.from_dict(

#                labels_map_new, orient='index', columns=['label_names'], dtype=str

#            )

            self.vae.init_latent_space(self.n_states, mu_new, log_pi_new)

            self.inferer = Inferer(self.n_states)

            self.mu = self.vae.latent_space.mu.numpy()

            self.pi = np.triu(np.ones(self.n_states))

            self.pi[self.pi > 0] = tf.nn.softmax(self.vae.latent_space.pi).numpy()[0]

    def train(self, stratify = False, test_size = 0.1, random_state: int = 0,

            learning_rate: float = 1e-3, batch_size: int = 256,

            L: int = 1, alpha: float = 0.10, beta: float = 1, gamma: float = 0, phi: float = 1,

            num_epoch: int = 200, num_step_per_epoch: Optional[int] =  None,

            early_stopping_patience: int = 10, early_stopping_tolerance: float = 0.01, 

            early_stopping_relative: bool = True, early_stopping_warmup: int = 0,

            path_to_weights: Optional[str] = None,

            verbose: bool = False, **kwargs):

        '''Train the model.

        Parameters

        ----------

        stratify : np.array, None, or False

            If an array is provided, or `stratify=None` and `self.labels` is available, then they will be used to perform stratified shuffle splitting. Otherwise, general shuffle splitting is used. Set to `False` if `self.labels` is not intended for stratified shuffle splitting.

        test_size : float or int, optional

            The proportion or size of the test set.

        random_state : int, optional

            The random state for data splitting.

        learning_rate : float, optional  

            The initial learning rate for the Adam optimizer.

        batch_size : int, optional  

            The batch size for training. Default is 256. Set to 32 if number of cells is small (less than 1000)

        L : int, optional  

            The number of MC samples.

        alpha : float, optional  

            The value of alpha in [0,1] to encourage covariate adjustment. Not used if there is no covariates.

        beta : float, optional  

            The value of beta in beta-VAE.

        gamma : float, optional

            The weight of mmd_loss.

        phi : float, optional

            The weight of Jacob norm of encoder.

        num_epoch : int, optional  

            The number of epoch.

        num_step_per_epoch : int, optional 

            The number of step per epoch, it will be inferred from number of cells and batch size if it is None.

        early_stopping_patience : int, optional 

            The maximum number of epochs if there is no improvement.

        early_stopping_tolerance : float, optional 

            The minimum change of loss to be considered as an improvement.

        early_stopping_relative : bool, optional

            Whether monitor the relative change of loss or not.            

        early_stopping_warmup : int, optional 

            The number of warmup epochs.            

        path_to_weights : str, optional 

            The path of weight file to be saved; not saving weight if None.

        **kwargs :  

            Extra key-value arguments for dimension reduction algorithms.        

        '''

        if gamma == 0 or self.conditions is None:

            conditions = np.array([np.nan] * self.adata.shape[0])

        else:

            conditions = self.conditions

        if stratify is None:

            stratify = self.labels

        elif stratify is False:

            stratify = None    

        id_train, id_test = train_test_split(

                                np.arange(self.X_input.shape[0]), 

                                test_size=test_size, 

                                stratify=stratify, 

                                random_state=random_state)

        if num_step_per_epoch is None:

            num_step_per_epoch = len(id_train)//batch_size+1

        c = None if self.covariates is None else self.covariates.astype(tf.keras.backend.floatx())

        self.train_dataset = train.warp_dataset(self.X_input[id_train].astype(tf.keras.backend.floatx()),

                                                None if c is None else c[id_train],

                                                batch_size, 

                                                self.X_output[id_train].astype(tf.keras.backend.floatx()), 

                                                self.scale_factor[id_train].astype(tf.keras.backend.floatx()),

                                                conditions = conditions[id_train],

                                                pi_cov = self.pi_cov[id_train])

        self.test_dataset = train.warp_dataset(self.X_input[id_test].astype(tf.keras.backend.floatx()),

                                                None if c is None else c[id_test],

                                                batch_size, 

                                                self.X_output[id_test].astype(tf.keras.backend.floatx()), 

                                                self.scale_factor[id_test].astype(tf.keras.backend.floatx()),

                                                conditions = conditions[id_test],

                                                pi_cov = self.pi_cov[id_test])

        self.vae = train.train(

            self.train_dataset,

            self.test_dataset,

            self.vae,

            learning_rate,

            L,

            alpha,

            beta,

            gamma,

            phi,

            num_epoch,

            num_step_per_epoch,

            early_stopping_patience,

            early_stopping_tolerance,

            early_stopping_relative,

            early_stopping_warmup,  

            verbose,

            **kwargs            

            )

        self.update_z()

        self.mu = self.vae.latent_space.mu.numpy()

        self.pi = np.triu(np.ones(self.n_states))

        self.pi[self.pi > 0] = tf.nn.softmax(self.vae.latent_space.pi).numpy()[0]

        if path_to_weights is not None:

            self.save_model(path_to_weights)

    def output_pi(self, pi_cov):

        """return a matrix n_states by n_states and a mask for plotting, which can be used to cover the lower triangular(except the diagnoals) of a heatmap"""

        p = self.vae.pilayer

        pi_cov = tf.expand_dims(tf.constant([pi_cov], dtype=tf.float32), 0)

        pi_val = tf.nn.softmax(p(pi_cov)).numpy()[0]

        # Create heatmap matrix

        n = self.vae.n_states

        matrix = np.zeros((n, n))

        matrix[np.triu_indices(n)] = pi_val

        mask = np.tril(np.ones_like(matrix), k=-1)

        return matrix, mask

    def return_pilayer_weights(self):

        """return parameters of pilayer, which has dimension dim(pi_cov) + 1 by n_categories, the last row is biases"""

        return np.vstack((model.vae.pilayer.weights[0].numpy(), model.vae.pilayer.weights[1].numpy().reshape(1, -1)))

    def posterior_estimation(self, batch_size: int = 32, L: int = 50, **kwargs):

        '''Initialize trajectory inference by computing the posterior estimations.        

        Parameters

        ----------

        batch_size : int, optional

            The batch size when doing inference.

        L : int, optional

            The number of MC samples when doing inference.

        **kwargs :  

            Extra key-value arguments for dimension reduction algorithms.              

        '''

        c = None if self.covariates is None else self.covariates.astype(tf.keras.backend.floatx())

        self.test_dataset = train.warp_dataset(self.X_input.astype(tf.keras.backend.floatx()), 

                                               c,

                                               batch_size)

        _, _, self.pc_x,\

            self.cell_position_posterior,self.cell_position_variance,_ = self.vae.inference(self.test_dataset, L=L)

        uni_cluster_labels = self.labels_map['label_names'].to_numpy()

        self.adata.obs['vitae_new_clustering'] = uni_cluster_labels[np.argmax(self.cell_position_posterior, 1)]

        self.adata.obs['vitae_new_clustering'] = self.adata.obs['vitae_new_clustering'].astype('category')

        print("New clustering labels saved as 'vitae_new_clustering' in self.adata.obs.")

        return None

    def infer_backbone(self, method: str = 'modified_map', thres = 0.5,

            no_loop: bool = True, cutoff: float = 0,

            visualize: bool = True, color = 'vitae_new_clustering',path_to_fig = None,**kwargs):

        ''' Compute edge scores.

        Parameters

        ----------

        method : string, optional

            'mean', 'modified_mean', 'map', or 'modified_map'.

        thres : float, optional

            The threshold used for filtering edges \(e_{ij}\) that \((n_{i}+n_{j}+e_{ij})/N<thres\), only applied to mean method.

        no_loop : boolean, optional

            Whether loops are allowed to exist in the graph. If no_loop is true, will prune the graph to contain only the

            maximum spanning true

        cutoff : string, optional

            The score threshold for filtering edges with scores less than cutoff.

        visualize: boolean

            whether plot the current trajectory backbone (undirected graph)

        Returns

        ----------

        G : nx.Graph

            The weighted graph with weight on each edge indicating its score of existence.

        '''

        # build_graph, return graph

        self.backbone = self.inferer.build_graphs(self.cell_position_posterior, self.pc_x,

                method, thres, no_loop, cutoff)

        self.cell_position_projected = self.inferer.modify_wtilde(self.cell_position_posterior, 

                np.array(list(self.backbone.edges)))

        uni_cluster_labels = self.labels_map['label_names'].to_numpy()

        temp_dict = {i:label for i,label in enumerate(uni_cluster_labels)}

        nx.relabel_nodes(self.backbone, temp_dict)

        self.adata.obs['vitae_new_clustering'] = uni_cluster_labels[np.argmax(self.cell_position_projected, 1)]

        self.adata.obs['vitae_new_clustering'] = self.adata.obs['vitae_new_clustering'].astype('category')

        print("'vitae_new_clustering' updated based on the projected cell positions.")

        self.uncertainty = np.sum((self.cell_position_projected - self.cell_position_posterior)**2, axis=-1) \

            + np.sum(self.cell_position_variance, axis=-1)

        self.adata.obs['projection_uncertainty'] = self.uncertainty

        print("Cell projection uncertainties stored as 'projection_uncertainty' in self.adata.obs")

        if visualize:

            self._adata.obs = self.adata.obs.copy()

            self.ax = self.plot_backbone(directed = False,color = color, **kwargs)

            if path_to_fig is not None:

                self.ax.figure.savefig(path_to_fig)

            self.ax.figure.show()

        return None

    def select_root(self, days, method: str = 'proportion'):

        '''Order the vertices/states based on cells' collection time information to select the root state.      

        Parameters

        ----------

        day : np.array 

            The day information for selected cells used to determine the root vertex.

            The dtype should be 'int' or 'float'.

        method : str, optional

            'sum' or 'mean'. 

            For 'proportion', the root is the one with maximal proportion of cells from the earliest day.

            For 'mean', the root is the one with earliest mean time among cells associated with it.

        Returns

        ----------

        root : int 

            The root vertex in the inferred trajectory based on given day information.

        '''

        ## TODO: change return description

        if days is not None and len(days)!=self.X_input.shape[0]:

            raise ValueError("The length of day information ({}) is not "

                "consistent with the number of selected cells ({})!".format(

                    len(days), self.X_input.shape[0]))

        if not hasattr(self, 'cell_position_projected'):

            raise ValueError("Need to call 'infer_backbone' first!")

        collection_time = np.dot(days, self.cell_position_projected)/np.sum(self.cell_position_projected, axis = 0)

        earliest_prop = np.dot(days==np.min(days), self.cell_position_projected)/np.sum(self.cell_position_projected, axis = 0)

        root_info = self.labels_map.copy()

        root_info['mean_collection_time'] = collection_time

        root_info['earliest_time_prop'] = earliest_prop

        root_info.sort_values('mean_collection_time', inplace=True)

        return root_info

    def plot_backbone(self, directed: bool = False, 

                      method: str = 'UMAP', color = 'vitae_new_clustering', **kwargs):

        '''Plot the current trajectory backbone (undirected graph).

        Parameters

        ----------

        directed : boolean, optional

            Whether the backbone is directed or not.

        method : str, optional

            The dimension reduction method to use. The default is "UMAP".

        color : str, optional

            The key for annotations of observations/cells or variables/genes, e.g., 'ann1' or ['ann1', 'ann2'].

            The default is 'vitae_new_clustering'.

        **kwargs :

            Extra key-value arguments that can be passed to scanpy plotting functions (scanpy.pl.XX).

        '''

        if not isinstance(color,str):

            raise ValueError('The color argument should be of type str!')

        ax = self.visualize_latent(method = method, color=color, show=False, **kwargs)

        dict_label_num = {j:i for i,j in self.labels_map['label_names'].to_dict().items()}

        uni_cluster_labels = self.adata.obs['vitae_init_clustering'].cat.categories

        cluster_labels = self.adata.obs['vitae_new_clustering'].to_numpy()

        embed_z = self._adata.obsm[self.dict_method_scname[method]]

        embed_mu = np.zeros((len(uni_cluster_labels), 2))

        for l in uni_cluster_labels:

            embed_mu[dict_label_num[l],:] = np.mean(embed_z[cluster_labels==l], axis=0)

        if directed:

            graph = self.directed_backbone

        else:

            graph = self.backbone

        edges = list(graph.edges)

        edge_scores = np.array([d['weight'] for (u,v,d) in graph.edges(data=True)])

        if max(edge_scores) - min(edge_scores) == 0:

            edge_scores = edge_scores/max(edge_scores)

        else:

            edge_scores = (edge_scores - min(edge_scores))/(max(edge_scores) - min(edge_scores))*3

        value_range = np.maximum(np.diff(ax.get_xlim())[0], np.diff(ax.get_ylim())[0])

        y_range = np.min(embed_z[:,1]), np.max(embed_z[:,1], axis=0)

        for i in range(len(edges)):

            points = embed_z[np.sum(self.cell_position_projected[:, edges[i]]>0, axis=-1)==2,:]

            points = points[points[:,0].argsort()]

            try:

                x_smooth, y_smooth = _get_smooth_curve(

                    points,

                    embed_mu[edges[i], :],

                    y_range

                    )

            except:

                x_smooth, y_smooth = embed_mu[edges[i], 0], embed_mu[edges[i], 1]

            ax.plot(x_smooth, y_smooth,

                '-',

                linewidth= 1 + edge_scores[i],

                color="black",

                alpha=0.8,

                path_effects=[pe.Stroke(linewidth=1+edge_scores[i]+1.5,

                                        foreground='white'), pe.Normal()],

                zorder=1

                )

            if directed:

                delta_x = embed_mu[edges[i][1], 0] - x_smooth[-2]

                delta_y = embed_mu[edges[i][1], 1] - y_smooth[-2]

                length = np.sqrt(delta_x**2 + delta_y**2) / 50 * value_range

                ax.arrow(

                        embed_mu[edges[i][1], 0]-delta_x/length,

                        embed_mu[edges[i][1], 1]-delta_y/length,

                        delta_x/length,

                        delta_y/length,

                        color='black', alpha=1.0,

                        shape='full', lw=0, length_includes_head=True,

                        head_width=np.maximum(0.01*(1 + edge_scores[i]), 0.03) * value_range,

                        zorder=2) 

        colors = self._adata.uns['vitae_new_clustering_colors']

        for i,l in enumerate(uni_cluster_labels):

            ax.scatter(*embed_mu[dict_label_num[l]:dict_label_num[l]+1,:].T, 

                       c=[colors[i]], edgecolors='white', # linewidths=10,  norm=norm,

                       s=250, marker='*', label=l)

        plt.setp(ax, xticks=[], yticks=[])

        box = ax.get_position()

        ax.set_position([box.x0, box.y0 + box.height * 0.1,

                            box.width, box.height * 0.9])

        if directed:

            ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05),

                fancybox=True, shadow=True, ncol=5)

        return ax

    def plot_center(self, color = "vitae_new_clustering", plot_legend = True, legend_add_index = True,

                    method: str = 'UMAP',ncol = 2,font_size = "medium",

                    add_egde = False, add_direct = False,**kwargs):

        '''Plot the center of each cluster in the latent space.

        Parameters

        ----------

        color : str, optional

            The color of the center of each cluster. Default is "vitae_new_clustering".

        plot_legend : bool, optional

            Whether to plot the legend. Default is True.

        legend_add_index : bool, optional

            Whether to add the index of each cluster in the legend. Default is True.

        method : str, optional

            The dimension reduction method used for visualization. Default is 'UMAP'.

        ncol : int, optional

            The number of columns in the legend. Default is 2.

        font_size : str, optional

            The font size of the legend. Default is "medium".

        add_egde : bool, optional

            Whether to add the edges between the centers of clusters. Default is False.

        add_direct : bool, optional

            Whether to add the direction of the edges. Default is False.

        '''

        if color not in ["vitae_new_clustering","vitae_init_clustering"]:

            raise ValueError("Can only plot center of vitae_new_clustering or vitae_init_clustering")

        dict_label_num = {j: i for i, j in self.labels_map['label_names'].to_dict().items()}

        if legend_add_index:

            self._adata.obs["index_"+color] = self._adata.obs[color].map(lambda x: dict_label_num[x])

            ax = self.visualize_latent(method=method, color="index_" + color, show=False, legend_loc="on data",

                                        legend_fontsize=font_size,**kwargs)

            colors = self._adata.uns["index_" + color + '_colors']

        else:

            ax = self.visualize_latent(method=method, color = color, show=False,**kwargs)

            colors = self._adata.uns[color + '_colors']

        uni_cluster_labels = self.adata.obs[color].cat.categories

        cluster_labels = self.adata.obs[color].to_numpy()

        embed_z = self._adata.obsm[self.dict_method_scname[method]]

        embed_mu = np.zeros((len(uni_cluster_labels), 2))

        for l in uni_cluster_labels:

            embed_mu[dict_label_num[l], :] = np.mean(embed_z[cluster_labels == l], axis=0)

        leg = (self.labels_map.index.astype(str) + " : " + self.labels_map.label_names).values

        for i, l in enumerate(uni_cluster_labels):

            ax.scatter(*embed_mu[dict_label_num[l]:dict_label_num[l] + 1, :].T,

                       c=[colors[i]], edgecolors='white', # linewidths=3,

                       s=250, marker='*', label=leg[i])

        if plot_legend:

            ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), ncol=ncol, markerscale=0.8, frameon=False)

        plt.setp(ax, xticks=[], yticks=[])

        box = ax.get_position()

        ax.set_position([box.x0, box.y0 + box.height * 0.1,

                         box.width, box.height * 0.9])

        if add_egde:

            if add_direct:

                graph = self.directed_backbone

            else:

                graph = self.backbone

            edges = list(graph.edges)

            edge_scores = np.array([d['weight'] for (u, v, d) in graph.edges(data=True)])

            if max(edge_scores) - min(edge_scores) == 0:

                edge_scores = edge_scores / max(edge_scores)

            else:

                edge_scores = (edge_scores - min(edge_scores)) / (max(edge_scores) - min(edge_scores)) * 3

            value_range = np.maximum(np.diff(ax.get_xlim())[0], np.diff(ax.get_ylim())[0])

            y_range = np.min(embed_z[:, 1]), np.max(embed_z[:, 1], axis=0)

            for i in range(len(edges)):

                points = embed_z[np.sum(self.cell_position_projected[:, edges[i]] > 0, axis=-1) == 2, :]

                points = points[points[:, 0].argsort()]

                try:

                    x_smooth, y_smooth = _get_smooth_curve(

                        points,

                        embed_mu[edges[i], :],

                        y_range

                    )

                except:

                    x_smooth, y_smooth = embed_mu[edges[i], 0], embed_mu[edges[i], 1]

                ax.plot(x_smooth, y_smooth,

                        '-',

                        linewidth=1 + edge_scores[i],

                        color="black",

                        alpha=0.8,

                        path_effects=[pe.Stroke(linewidth=1 + edge_scores[i] + 1.5,

                                                foreground='white'), pe.Normal()],

                        zorder=1

                        )

                if add_direct:

                    delta_x = embed_mu[edges[i][1], 0] - x_smooth[-2]

                    delta_y = embed_mu[edges[i][1], 1] - y_smooth[-2]

                    length = np.sqrt(delta_x ** 2 + delta_y ** 2) / 50 * value_range

                    ax.arrow(

                        embed_mu[edges[i][1], 0] - delta_x / length,

                        embed_mu[edges[i][1], 1] - delta_y / length,

                        delta_x / length,

                        delta_y / length,

                        color='black', alpha=1.0,

                        shape='full', lw=0, length_includes_head=True,

                        head_width=np.maximum(0.01 * (1 + edge_scores[i]), 0.03) * value_range,

                        zorder=2)

        self.ax = ax

        self.ax.figure.show()

        return None

    def infer_trajectory(self, root: Union[int,str], digraph = None, color = "pseudotime",

                         visualize: bool = True, path_to_fig = None,  **kwargs):

        '''Infer the trajectory.

        Parameters

        ----------

        root : int or string

            The root of the inferred trajectory. Can provide either an int (vertex index) or string (label name)

        digraph : nx.DiGraph, optional

            The directed graph to be used for trajectory inference. If None, the minimum spanning tree of the estimated trajectory backbone will be used.

        cutoff : string, optional

            The threshold for filtering edges with scores less than cutoff.

        visualize: boolean

            Whether plot the current trajectory backbone (directed graph)

        path_to_fig : string, optional  

            The path to save figure, or don't save if it is None.

        **kwargs : dict, optional

            Other keywords arguments for plotting.

        '''

        if isinstance(root,str):

            if root not in self.labels_map.values:

                raise ValueError("Root {} is not in the label names!".format(root))

            root = self.labels_map[self.labels_map['label_names']==root].index[0]

        if digraph is None:

            connected_comps = nx.node_connected_component(self.backbone, root)

            subG = self.backbone.subgraph(connected_comps)

            ## generate directed backbone which contains no loops

            DG = nx.DiGraph(nx.to_directed(self.backbone))

            temp = DG.subgraph(connected_comps)

            DG.remove_edges_from(temp.edges - nx.dfs_edges(DG, root))

            self.directed_backbone = DG

        else:

            if not nx.is_directed_acyclic_graph(digraph):

                raise ValueError("The graph 'digraph' should be a directed acyclic graph.")

            if set(digraph.nodes) != set(self.backbone.nodes):

                raise ValueError("The nodes in 'digraph' do not match the nodes in 'self.backbone'.")

            self.directed_backbone = digraph

            connected_comps = nx.node_connected_component(digraph, root)

            subG = self.backbone.subgraph(connected_comps)

        if len(subG.edges)>0:

            milestone_net = self.inferer.build_milestone_net(subG, root)

            if self.inferer.no_loop is False and milestone_net.shape[0]<len(self.backbone.edges):

                warnings.warn("The directed graph shown is a minimum spanning tree of the estimated trajectory backbone to avoid arbitrary assignment of the directions.")

            self.pseudotime = self.inferer.comp_pseudotime(milestone_net, root, self.cell_position_projected)

        else:

            warnings.warn("There are no connected states for starting from the giving root.")

            self.pseudotime = -np.ones(self._adata.shape[0])

        self.adata.obs['pseudotime'] = self.pseudotime

        print("Cell projection uncertainties stored as 'pseudotime' in self.adata.obs")

        if visualize:

            self._adata.obs['pseudotime'] = self.pseudotime

            self.ax = self.plot_backbone(directed = True, color = color, **kwargs)

            if path_to_fig is not None:

                self.ax.figure.savefig(path_to_fig)

            self.ax.figure.show()

        return None

    def differential_expression_test(self, alpha: float = 0.05, cell_subset = None, order: int = 1):

        '''Differentially gene expression test. All (selected and unselected) genes will be tested 

        Only cells in `selected_cell_subset` will be used, which is useful when one need to

        test differentially expressed genes on a branch of the inferred trajectory.

        Parameters

        ----------

        alpha : float, optional

            The cutoff of p-values.

        cell_subset : np.array, optional

            The subset of cells to be used for testing. If None, all cells will be used.

        order : int, optional

            The maxium order we used for pseudotime in regression.

        Returns

        ----------

        res_df : pandas.DataFrame

            The test results of expressed genes with two columns,

            the estimated coefficients and the adjusted p-values.

        '''

        if not hasattr(self, 'pseudotime'):

            raise ReferenceError("Pseudotime does not exist! Please run 'infer_trajectory' first.")

        if cell_subset is None:

            cell_subset = np.arange(self.X_input.shape[0])

            print("All cells are selected.")

        if order < 1:

            raise  ValueError("Maximal order of pseudotime in regression must be at least 1.")

        # Prepare X and Y for regression expression ~ rank(PDT) + covariates

        Y = self.adata.X[cell_subset,:]

#        std_Y = np.std(Y, ddof=1, axis=0, keepdims=True)

#        Y = np.divide(Y-np.mean(Y, axis=0, keepdims=True), std_Y, out=np.empty_like(Y)*np.nan, where=std_Y!=0)

        X = stats.rankdata(self.pseudotime[cell_subset])        

        if order > 1:

            for _order in range(2, order+1):

                X = np.c_[X, X**_order]

        X = ((X-np.mean(X,axis=0, keepdims=True))/np.std(X, ddof=1, axis=0, keepdims=True))

        X = np.c_[np.ones((X.shape[0],1)), X]

        if self.covariates is not None:

            X = np.c_[X, self.covariates[cell_subset, :]]

        res_df = DE_test(Y, X, self.adata.var_names, i_test = np.array(list(range(1,order+1))), alpha = alpha)

        return res_df[res_df.pvalue_adjusted_1 != 0]

    def evaluate(self, milestone_net, begin_node_true, grouping = None,

                thres: float = 0.5, no_loop: bool = True, cutoff: Optional[float] = None,

                method: str = 'mean', path: Optional[str] = None):

        ''' Evaluate the model.

        Parameters

        ----------

        milestone_net : pd.DataFrame

            The true milestone network. For real data, milestone_net will be a DataFrame of the graph of nodes.

            Eg.

            from|to

            ---|---

            cluster 1 | cluster 1

            cluster 1 | cluster 2

            For synthetic data, milestone_net will be a DataFrame of the (projected)

            positions of cells. The indexes are the orders of cells in the dataset.

            Eg.

            from|to|w

            ---|---|---

            cluster 1 | cluster 1 | 1

            cluster 1 | cluster 2 | 0.1

        begin_node_true : str or int

            The true begin node of the milestone.

        grouping : np.array, optional

            \([N,]\) The labels. For real data, grouping must be provided.

        Returns

        ----------

        res : pd.DataFrame

            The evaluation result.

        '''

        if not hasattr(self, 'labels_map'):

            raise ValueError("No given labels for training.")

        '''

        # Evaluate for the whole dataset will ignore selected_cell_subset.

        if len(self.selected_cell_subset)!=len(self.cell_names):

            warnings.warn("Evaluate for the whole dataset.")

        '''

        # If the begin_node_true, need to encode it by self.le.

        # this dict is for milestone net cause their labels are not merged

        # all keys of label_map_dict are str

        label_map_dict = dict()

        for i in range(self.labels_map.shape[0]):

            label_mapped = self.labels_map.loc[i]

            ## merged cluster index is connected by comma

            for each in label_mapped.values[0].split(","):

                label_map_dict[each] = i

        if isinstance(begin_node_true, str):

            begin_node_true = label_map_dict[begin_node_true]