import pandas as pd
import sklearn.neighbors
import scipy.sparse as sp
import networkx as nx
import numpy as np
import torch
import argparse

def parameter_setting():

    parser = argparse.ArgumentParser(description='stClinic for predicting clinically-related niches')
    parser.add_argument('--input_dir',   '-IP', type = str, default = '/home/zuocm/Share_data/xiajunjie/stClinic/Datasets/DLPFC/',  help='data directory')
    parser.add_argument('--out_dir',   '-OP', type = str, default = '/home/zuocm/Share_data/xiajunjie/stClinic/Datasets/DLPFC/stClinic/',  help='data directory')
    parser.add_argument('--rad_cutoff',   '-RC', type = int, default = 150,    help='The radius value used to construct intra-edge (i.e., spatial nearest neighbors) for each slice')
    parser.add_argument('--k_cutoff',   '-KC', type = int, default = 6,    help='The number of spatial neighbors used to construct intra-edge')
    parser.add_argument('--k',   '-K', type = int, default = 5,    help='The number of mutual nearest neighbors used to construct inter-edges across slices')
    parser.add_argument('--n_top_genes',   '-NTG', type = int, default = 5000,    help='The number of highly variable genes selected for each slice')
    parser.add_argument('--n_centroids',   '-NCS', type = int, default = 7,    help='The number of components of the GMM')
    parser.add_argument('--lr_integration',   '-LRI', type = float, default = 0.0005,    help='The learning rate used by stClinic when extracting batch-corrected features in slices')
    parser.add_argument('--lr_prediction',   '-LRP', type = float, default = 0.05,    help='The learning rate used by supervised stClinic')


def Cal_Spatial_Net(adata, rad_cutoff=None, k_cutoff=None,
                    max_neigh=50, model='Radius', verbose=True):
    """\
    Construct the spatial neighbor networks.

    Parameters
    ----------
    adata
        AnnData object of scanpy package.
    rad_cutoff
        radius cutoff when model='Radius'
    k_cutoff
        The number of nearest neighbors when model='KNN'
    model
        The network construction model. When model=='Radius', the spot is connected to spots whose distance is less than rad_cutoff. When model=='KNN', the spot is connected to its first k_cutoff nearest neighbors.

    Returns
    -------
    The spatial networks are saved in adata.uns['Spatial_Net']
    """

    assert (model in ['Radius', 'KNN'])
    if verbose:
        print('------Calculating intra-spatial graph...')
    coor = pd.DataFrame(adata.obsm['spatial'])
    coor.index = adata.obs.index
    coor.columns = ['imagerow', 'imagecol']

    nbrs = sklearn.neighbors.NearestNeighbors(
        n_neighbors=max_neigh + 1, algorithm='ball_tree').fit(coor)
    distances, indices = nbrs.kneighbors(coor)
    if model == 'KNN':
        indices = indices[:, 1:k_cutoff + 1]
        distances = distances[:, 1:k_cutoff + 1]
    if model == 'Radius':
        indices = indices[:, 1:]
        distances = distances[:, 1:]

    KNN_list = []
    for it in range(indices.shape[0]):
        KNN_list.append(pd.DataFrame(zip([it] * indices.shape[1], indices[it, :], distances[it, :])))
    KNN_df = pd.concat(KNN_list)
    KNN_df.columns = ['Cell1', 'Cell2', 'Distance']

    Spatial_Net = KNN_df.copy()
    if model == 'Radius':
        Spatial_Net = KNN_df.loc[KNN_df['Distance'] < rad_cutoff,]
    id_cell_trans = dict(zip(range(coor.shape[0]), np.array(coor.index), ))
    Spatial_Net['Cell1'] = Spatial_Net['Cell1'].map(id_cell_trans)
    Spatial_Net['Cell2'] = Spatial_Net['Cell2'].map(id_cell_trans)
    # self_loops = pd.DataFrame(zip(Spatial_Net['Cell1'].unique(), Spatial_Net['Cell1'].unique(),
    #                  [0] * len((Spatial_Net['Cell1'].unique())))) ###add self loops
    # self_loops.columns = ['Cell1', 'Cell2', 'Distance']
    # Spatial_Net = pd.concat([Spatial_Net, self_loops], axis=0)

    if verbose:
        print('The graph contains %d edges, %d cells.' % (Spatial_Net.shape[0], adata.n_obs))
        print('%.4f neighbors per spot on average.' % (Spatial_Net.shape[0] / adata.n_obs))
    adata.uns['Spatial_Net'] = Spatial_Net

    #########
    X = pd.DataFrame(adata.X.toarray()[:, ], index=adata.obs.index, columns=adata.var.index)
    cells = np.array(X.index)
    cells_id_tran = dict(zip(cells, range(cells.shape[0])))
    if 'Spatial_Net' not in adata.uns.keys():
        raise ValueError("Spatial_Net is not existed! Run Cal_Spatial_Net first!")
        
    Spatial_Net = adata.uns['Spatial_Net']
    G_df = Spatial_Net.copy()
    G_df['Cell1'] = G_df['Cell1'].map(cells_id_tran)
    G_df['Cell2'] = G_df['Cell2'].map(cells_id_tran)
    G = sp.coo_matrix((np.ones(G_df.shape[0]), (G_df['Cell1'], G_df['Cell2'])), shape=(adata.n_obs, adata.n_obs))
    G = G + sp.eye(G.shape[0])  # self-loop
    adata.uns['adj'] = G



def mclust_R(adata, num_cluster, modelNames='EEE', used_obsm='VGAEX', random_seed=666):
    """\
    Clustering using the mclust algorithm.
    The parameters are the same as those in the R package mclust.
    """

    np.random.seed(random_seed)
    import rpy2.robjects as robjects
    robjects.r.library("mclust")

    import rpy2.robjects.numpy2ri
    rpy2.robjects.numpy2ri.activate()
    r_random_seed = robjects.r['set.seed']
    r_random_seed(random_seed)
    rmclust = robjects.r['Mclust']

    mclust_res = np.array(rmclust(np.array(adata.obsm[used_obsm]), num_cluster, modelNames)[-2])

    adata.obs['mclust'] = mclust_res
    adata.obs['mclust'] = adata.obs['mclust'].astype('int')
    adata.obs['mclust'] = adata.obs['mclust'].astype('category')
    return adata



def match_cluster_labels(true_labels,est_labels):
    true_labels_arr = np.array(list(true_labels))
    est_labels_arr = np.array(list(est_labels))
    org_cat = list(np.sort(list(pd.unique(true_labels))))
    est_cat = list(np.sort(list(pd.unique(est_labels))))
    B = nx.Graph()
    B.add_nodes_from([i+1 for i in range(len(org_cat))], bipartite=0)
    B.add_nodes_from([-j-1 for j in range(len(est_cat))], bipartite=1)
    for i in range(len(org_cat)):
        for j in range(len(est_cat)):
            weight = np.sum((true_labels_arr==org_cat[i])* (est_labels_arr==est_cat[j]))
            B.add_edge(i+1,-j-1, weight=-weight)
    match = nx.algorithms.bipartite.matching.minimum_weight_full_matching(B)
#     match = minimum_weight_full_matching(B)
    if len(org_cat)>=len(est_cat):
        return np.array([match[-est_cat.index(c)-1]-1 for c in est_labels_arr])
    else:
        unmatched = [c for c in est_cat if not (-est_cat.index(c)-1) in match.keys()]
        l = []
        for c in est_labels_arr:
            if (-est_cat.index(c)-1) in match: 
                l.append(match[-est_cat.index(c)-1]-1)
            else:
                l.append(len(org_cat)+unmatched.index(c))
        return np.array(l)  



def count_params(model):
    all_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    non_trainable_params = sum(p.numel() for p in model.parameters() if not p.requires_grad)

    print(f'Total params: {all_params}\nTrainable params: {trainable_params}\nNon-trainable params: {non_trainable_params}')



from sklearn.preprocessing import scale
from scipy.sparse.linalg import eigsh
def estimate_k(data):
    """
    Estimate number of groups k:
        based on random matrix theory (RTM), borrowed from SC3
        input data is (p,n) matrix, p is feature, n is sample
    """
    p, n = data.shape

    x = scale(data, with_mean=False)
    muTW = (np.sqrt(n-1) + np.sqrt(p)) ** 2
    sigmaTW = (np.sqrt(n-1) + np.sqrt(p)) * (1/np.sqrt(n-1) + 1/np.sqrt(p)) ** (1/3)
    sigmaHatNaive = x.T.dot(x)

    bd = np.sqrt(p) * sigmaTW + muTW
    evals, _ = eigsh(sigmaHatNaive)

    k = 0
    for i in range(len(evals)):
        if evals[i] > bd:
            k += 1
    return k



from sklearn.metrics import silhouette_score as s_score
def F1score(adata):
    """
    Compute silhouette coefficient:
    F1score = 2*(1-silh'batch)*silh'cluster / [silh'cluster + (1-silh'batch)]
    silh'batch = (1+silh_batch)/2
    silh'cluster = (1+silh_cluster)/2
    """
    s_cluster = s_score(adata.obsm['X_umap'], adata.obs['mclust'])
    s_batch = s_score(adata.obsm['X_umap'], adata.obs['batch_name'])

    s_cluster = (1+s_cluster)/2
    s_batch = (1+s_batch)/2

    F1score = 2*(1-s_batch)*s_cluster / (s_cluster + (1-s_batch))

    return F1score



def LISI_score(adata, random_seed=666):
    """
    Compute cell-type / integration local inverse simpson'index
    """

    np.random.seed(random_seed)
    import rpy2.robjects as robjects
    robjects.r.library("lisi")
    rlisi = robjects.r['compute_lisi']

    import rpy2.robjects.numpy2ri, rpy2.robjects.pandas2ri
    rpy2.robjects.numpy2ri.activate()
    rpy2.robjects.pandas2ri.activate()

    r_random_seed = robjects.r['set.seed']
    r_random_seed(random_seed)

    robjects.globalenv['X'] = np.array(adata.obsm['X_umap'])
    meta_data = adata.obs[['Ground Truth', 'batch_name']]
    meta_data.columns = ['label1','label2']
    robjects.globalenv['meta_data'] = meta_data
    robjects.globalenv['cnames'] = robjects.r['c']('label1', 'label2')

    lisi_res = np.array(
                        rlisi(robjects.globalenv['X'], robjects.globalenv['meta_data'], robjects.globalenv['cnames'])
                        )

    clisi = np.mean(lisi_res, axis=1)[0]
    ilisi = np.mean(lisi_res, axis=1)[1]

    return dict(cLISI = clisi, iLISI = ilisi, LISI=lisi_res)



from lifelines.utils import concordance_index
from lifelines.statistics import logrank_test
def CIndex_lifeline(hazards, labels, survtime_all):
    return(concordance_index(survtime_all, -hazards, labels))



def cox_log_rank(hazardsdata, labels, survtime_all):
    median = np.median(hazardsdata)
    hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)
    hazards_dichotomize[hazardsdata > median] = 1
    idx = hazards_dichotomize == 0
    T1 = survtime_all[idx]
    T2 = survtime_all[~idx]
    E1 = labels[idx]
    E2 = labels[~idx]
    results = logrank_test(T1, T2, event_observed_A=E1, event_observed_B=E2)
    pvalue_pred = results.p_value
    return(pvalue_pred)



def accuracy_cox(hazardsdata, labels):
    # This accuracy is based on estimated survival events against true survival events
    median = np.median(hazardsdata)
    hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)
    hazards_dichotomize[hazardsdata > median] = 1
    correct = np.sum(hazards_dichotomize == labels)
    return correct / len(labels)



from lifelines import KaplanMeierFitter
from lifelines import CoxPHFitter
def KM_plot(hazardsdata, labels, survtime_all, output_dir):
    median = np.median(hazardsdata)
    hazards_dichotomize = np.zeros([len(hazardsdata)], dtype=int)
    hazards_dichotomize[hazardsdata > median] = 1
    idx = (hazards_dichotomize==1)

    T = survtime_all
    E = labels

    import matplotlib.pyplot as plt
    plt.figure(figsize=(8, 8))
    ax = plt.subplot(111)

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_linewidth(1)
    ax.spines['bottom'].set_linewidth(1)

    kmf_high = KaplanMeierFitter()
    ax = kmf_high.fit(T[idx], E[idx], label='high risk (n={})'.format(sum(idx))).plot_survival_function(linewidth=3, show_censors=True, censor_styles={'ms':9, 'marker':'+'}, ax=ax, ci_show=False, color='#FF4B68')

    kmf_low = KaplanMeierFitter()
    ax = kmf_low.fit(T[~idx], E[~idx], label='low risk (n={})'.format(len(idx)-sum(idx))).plot_survival_function(linewidth=3, show_censors=True, censor_styles={'ms':9, 'marker':'+'}, ax=ax, ci_show=False, color='#118DF0')

    legend = plt.legend(loc="upper right",fontsize=10)
    legend.get_frame().set_facecolor('none')
    legend.get_frame().set_linewidth(0.0)
    plt.xlabel('Time (days)',fontsize=12)
    plt.ylim([-0.01, 1.01])
    plt.ylabel('Population at risk (%)',fontsize=12)
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)

    logrank_stat = logrank_test(T[idx], T[~idx], E[idx], E[~idx], alpha=.99).p_value
    plt.text(7.5,0.05,s='p={:.2e}'.format(logrank_stat),fontsize=10)

    cph = CoxPHFitter()
    df = pd.DataFrame({'Surv_time': T, 'Event':T, 'High_risk':hazards_dichotomize})
    cph.fit(df, 'Surv_time', 'Event')
    hr_value = cph.summary.loc['High_risk', 'exp(coef)']
    ci_lower = cph.summary.loc['High_risk', 'exp(coef) lower 95%']
    ci_upper = cph.summary.loc['High_risk', 'exp(coef) upper 95%']

    plt.text(7.5,0.07, s=f'HR: {hr_value:.3f} ({ci_lower:.3f}, {ci_upper:.3f})', fontsize=10)

    plt.savefig(output_dir + 'Stat_AttPred_UMAP_KM_plot.jpg', dpi=500)



from sklearn.metrics import confusion_matrix
def CM_plot(pred_data, labels, output_dir):
    cm = confusion_matrix(labels, pred_data)

    import matplotlib.pyplot as plt
    import seaborn as sns
    plt.figure(figsize=(6,6))
    sns.heatmap(cm, annot=True, cmap='OrRd', fmt='g', annot_kws={'fontsize': 12})
    plt.xlabel('Ground-truth label', fontdict={'size':12})
    plt.ylabel('Predicted label', fontdict={'size':12})

    plt.savefig(output_dir + 'Stat_AttPred_UMAP_CM_plot.jpg', dpi=500)



from sklearn.metrics import roc_curve, auc
def ROC_plot(pred_data, labels, output_dir):
    fpr, tpr, thresholds = roc_curve(labels, pred_data)
    auc_score = auc(fpr,tpr)

    import matplotlib.pyplot as plt
    plt.figure(figsize=(6,6))
    plt.plot(fpr, tpr, label=f'AUC: {auc_score:.3f}')
    plt.plot([0,1],[0,1],linestyle='--',color='grey')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.0])
    plt.title('ROC curve')
    plt.xlabel('False Positive Rate',fontdict={'size':12})
    plt.ylabel('True Positive Rate',fontdict={'size':12})
    plt.legend(fontsize=12,loc='lower right')

    plt.savefig(output_dir + 'Stat_AttPred_UMAP_ROC_plot.jpg', dpi=500)



from sklearn.metrics import pairwise_distances
def knn_smoothing(latent, k, mat):
    dist = pairwise_distances(latent)
    row = []
    col = []
    sorted_knn = dist.argsort(axis=1)
    for idx in list(range(np.shape(dist)[0])):
        col.extend(sorted_knn[idx, : k].tolist())
        row.extend([idx] * k)

    res = np.zeros((mat.shape[0], mat.shape[1]))
    for i in range(len(col)):
        res[row[i]] += mat[col[i]]

    return res



from typing import List, Mapping, Optional, Union
import faiss
from sklearn.metrics.pairwise import euclidean_distances
from anndata import AnnData
def spatial_match(embds:List[torch.Tensor],
                  reorder:Optional[bool]=True,
                  smooth:Optional[bool]=True,
                  smooth_range:Optional[int]=20,
                  scale_coord:Optional[bool]=True,
                  adatas:Optional[List[AnnData]]=None,
                  return_euclid:Optional[bool]=False,
                  verbose:Optional[bool]=False,
                  get_null_distri:Optional[bool]=False
    )-> List[Union[np.ndarray,torch.Tensor]]:
    r"""
    Use embedding to match cells from different batches based on cosine similarity
    
    Parameters
    ----------
    embds
        list of embeddings
    reorder
        if reorder embedding by cell numbers
    smooth
        if smooth the mapping by Euclid distance
    smooth_range
        use how many candidates to do smooth
    scale_coord
        if scale the coordinate to [0,1]
    adatas
        list of adata object
    verbose
        if print log
    get_null_distri
        if get null distribution of cosine similarity
    
    Note
    ----------
    Automatically use larger dataset as source
    
    Return
    ----------
    Best matching, Top n matching and cosine similarity matrix of top n  
    """
    if reorder and embds[0].shape[0] < embds[1].shape[0]:
        embd0 = embds[1]
        embd1 = embds[0]
        adatas = adatas[::-1] if adatas is not None else None
    else:
        embd0 = embds[0]
        embd1 = embds[1]
        
    if get_null_distri:
        embd0 = torch.tensor(embd0)
        embd1 = torch.tensor(embd1)
        sample1_index = torch.randint(0, embd0.shape[0], (1000,))
        sample2_index = torch.randint(0, embd1.shape[0], (1000,))
        cos = torch.nn.CosineSimilarity(dim=1)
        null_distri = cos(embd0[sample1_index], embd1[sample2_index])

    index = faiss.index_factory(embd1.shape[1], "Flat", faiss.METRIC_INNER_PRODUCT)
    embd0_np = embd0.detach().cpu().numpy() if torch.is_tensor(embd0) else embd0
    embd1_np = embd1.detach().cpu().numpy() if torch.is_tensor(embd1) else embd1
    embd0_np = embd0_np.copy().astype('float32')
    embd1_np = embd1_np.copy().astype('float32')
    faiss.normalize_L2(embd0_np)
    faiss.normalize_L2(embd1_np)
    index.add(embd0_np)
    similarity, order = index.search(embd1_np, smooth_range)
    best = []
    if smooth and adatas != None:
        if verbose:
            print('Smoothing mapping, make sure object is in same direction')
        if scale_coord:
            # scale spatial coordinate of every adata to [0,1]
            adata1_coord = adatas[0].obsm['spatial'].copy()
            adata2_coord = adatas[1].obsm['spatial'].copy()
            for i in range(2):
                    adata1_coord[:,i] = (adata1_coord[:,i]-np.min(adata1_coord[:,i]))/(np.max(adata1_coord[:,i])-np.min(adata1_coord[:,i]))
                    adata2_coord[:,i] = (adata2_coord[:,i]-np.min(adata2_coord[:,i]))/(np.max(adata2_coord[:,i])-np.min(adata2_coord[:,i]))
        dis_list = []
        for query in range(embd1_np.shape[0]):
            ref_list = order[query, :smooth_range]
            dis = euclidean_distances(adata2_coord[query,:].reshape(1, -1),
                                      adata1_coord[ref_list,:])
            dis_list.append(dis)
            best.append(ref_list[np.argmin(dis)])
    else:
        best = order[:,0]

    if return_euclid and smooth and adatas != None:
        dis_array = np.squeeze(np.array(dis_list))
        if get_null_distri:
            return np.array(best), order, similarity, dis_array, null_distri
        else:
            return np.array(best), order, similarity, dis_array
    else:
        return np.array(best), order, similarity



class EarlyStopping:
    """
    Early stops the training if loss doesn't improve after a given patience.
    """
    def __init__(self, patience=10, verbose=False, checkpoint_file=''):
        """
        Parameters
        ----------
        patience 
            How long to wait after last time loss improved. Default: 10
        verbose
            If True, prints a message for each loss improvement. Default: False
        """
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.loss_min = np.Inf
        self.checkpoint_file = checkpoint_file

    def __call__(self, loss, model):
        if np.isnan(loss):
            self.early_stop = True
        score = -loss

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(loss, model)
        elif score <= self.best_score:
            self.counter += 1
            if self.verbose:
                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
                model.load_model(self.checkpoint_file)
        else:
            self.best_score = score
            self.save_checkpoint(loss, model)
            self.counter = 0

    def save_checkpoint(self, loss, model):
        '''
        Saves model when loss decrease.
        '''
        if self.verbose:
            print(f'Loss decreased ({self.loss_min:.6f} --> {loss:.6f}).  Saving model ...')
        if self.checkpoint_file:
            torch.save(model.state_dict(), self.checkpoint_file)
        self.loss_min = loss