# -*- coding: utf-8 -*-
# @Author: dongqing
# @Date:   2024-01-07 19:26:42
# @Last Modified by:   dongqing
# @Last Modified time: 2024-04-17 03:43:30

import os, sys
import argparse
import functools
import pandas as pd
import numpy as np
import multiprocessing as mp
import scanpy as sc
import anndata
import concurrent.futures

from sklearn_ann.kneighbors.annoy import AnnoyTransformer
from Cellist.IO import read_10X_h5, write_10X_h5
from annoy import AnnoyIndex
from scipy.sparse import csr_matrix, csc_matrix, hstack

from Cellist.Utility import *

def ImputationParser(subparsers):
    parser = subparsers.add_parser("impute", 
        help = "Perform spatially-aware gene imputation within each cluster. ")
    group_input = parser.add_argument_group("Input arguments")
    group_input.add_argument("--expr", dest = "expr", help = "Cell-level gene expression file in h5 format generated by 'Cellist seg'. ")
    group_input.add_argument("--spatial", dest = "spatial", help = "Spatial coordniates of cells generated by 'Cellist seg'. ")
    group_input.add_argument("--nworkers", dest = "num_workers", type = int, default = 8,
    help = "Maximum number of workers to use. ")

    group_output = parser.add_argument_group("Output arguments")
    group_output.add_argument("--outdir", dest = "out_dir",default = "",help = "Output directory.")
    group_output.add_argument("--outprefix", dest = "out_prefix",default = "",help = "Output prefix.")
    return(parser)

def ST_clustering(st_count_mat, st_count_genes, st_count_cells, spatial_df, out_dir, out_prefix):
    adata = anndata.AnnData(st_count_mat.transpose(), obs = dict(obs_names = st_count_cells), var = dict(var_names = st_count_genes))
    # preprocess
    sc.pp.filter_cells(adata, min_genes = 50)
    sc.pp.filter_genes(adata, min_cells = 100)
    sc.pp.calculate_qc_metrics(adata, inplace = True)
    # normalize data
    sc.pp.normalize_total(adata, target_sum=1e4)
    sc.pp.log1p(adata)
    sc.pp.highly_variable_genes(adata, flavor='seurat_v3', n_top_genes = 2000)
    sc.pp.scale(adata)
    sc.pp.pca(adata)
    sc.pp.neighbors(adata, n_pcs = 20, transformer=AnnoyTransformer(n_neighbors=15, n_trees=50))
    sc.tl.umap(adata)
    res = 0.6
    sc.tl.louvain(adata, resolution = res, key_added = 'louvain_%s' %res)
    # sc.tl.leiden(adata, n_iterations=2)
    fig, ax = plt.subplots(figsize = (6,6))
    sc.pl.umap(adata, color = 'louvain_%s' %res, ax = ax)
    plt.title('')
    plt.savefig(os.path.join(out_dir, "%s_louvain_UMAP.pdf" %out_prefix), bbox_inches='tight')
    # add celltype info
    adata.obsm["spatial"] = spatial_df.loc[adata.obs.index, ["y", "x"]].to_numpy()
    fig, ax = plt.subplots(figsize = (12,12))
    if adata.obs.shape[0] < 50000:
        scatter_size = 12
    else:
        scatter_size = 6
    sc.pl.embedding(adata, basis="spatial", color = 'louvain_%s' %res, s=scatter_size, show=False, ax = ax)
    ax.set_aspect('equal', adjustable='box')
    ax.invert_yaxis()
    plt.title('')
    plt.axis('off')
    plt.savefig(os.path.join(out_dir, "%s_louvain_spatial.pdf" %out_prefix), bbox_inches='tight')
    return(adata)

def CalAffinity_PCA(PCA_tree, i, k_params):
    nn, nd = PCA_tree.get_nns_by_item(i = i, n = k_params, include_distances=True)
    row_list = [i]*k_params
    col_list = nn
    dist_list = nd
    sigma = nd[int(k_params/3)]
    sigma_list = [sigma]*k_params
    return(np.array([row_list, col_list, dist_list, sigma_list]))

def CalAffinity_Spatial(Spatial_tree, i, k_params):
    nn, nd = Spatial_tree.get_nns_by_item(i = i, n = k_params, include_distances=True)
    row_list = [i]*k_params
    col_list = nn
    dist_list = nd
    sigma = nd[int(k_params/3)]
    sigma_list = [sigma]*k_params
    return(np.array([row_list, col_list, dist_list, sigma_list]))

def BuildTree(data, n_trees = 50):
    f = data.shape[1]
    t = AnnoyIndex(f, 'euclidean')
    n_samples = data.shape[0]
    for i in range(n_samples):
        v = data.iloc[i, :].tolist()
        t.add_item(i, v)
    t.build(n_trees)
    return(t)

def CalAffinityMat(result_list, n_samples):
    res_array = np.hstack(result_list)
    aff_array = np.exp(-np.square(res_array[2,:]/res_array[3,:]))
    aff_mat = csr_matrix((aff_array, (res_array[0,:], res_array[1,:])), shape=(n_samples, n_samples))
    aff_mat_sym = (aff_mat + aff_mat.T)/2
    return(aff_mat_sym)

def CalCombineAffinityMat(result_list, n_samples):
    res_array = np.hstack(result_list)
    aff_mat = csr_matrix((res_array[2,:], (res_array[0,:], res_array[1,:])), shape=(n_samples, n_samples))
    return(aff_mat)

def CombineAffinityMat(pca_aff_mat_sym, spatial_aff_mat_sym, knn, i):
    pca_nn = np.where(pca_aff_mat_sym[i,:].toarray()[0] > 0)[0]
    spatial_nn = np.where(spatial_aff_mat_sym[i,:].toarray()[0] > 0)[0]
    common_nn = np.intersect1d(pca_nn, spatial_nn)
    if common_nn.shape[0] > 5:
        if common_nn.shape[0] > knn:
            pca_aff_value = pca_aff_mat_sym[i, common_nn].toarray()[0]
            cut_off = np.quantile(pca_aff_value, (pca_aff_value.shape[0] - knn - 1)/pca_aff_value.shape[0])
            common_nn_cut = common_nn[pca_aff_value > cut_off]
        else:
            common_nn_cut = common_nn
        self_ind = np.where(common_nn_cut == i)[0]
        spatial_aff_array = spatial_aff_mat_sym[i, common_nn_cut].toarray()[0]
        spatial_aff_array[np.where(common_nn_cut != i)[0]] = spatial_aff_array[np.where(common_nn_cut != i)[0]]/np.sum(spatial_aff_array[np.where(common_nn_cut != i)[0]])
        spatial_aff_array[self_ind] = 1
        pca_aff_array = pca_aff_mat_sym[i, common_nn_cut].toarray()[0]
        pca_aff_array[np.where(common_nn_cut != i)[0]] = pca_aff_array[np.where(common_nn_cut != i)[0]]/np.sum(pca_aff_array[np.where(common_nn_cut != i)[0]])
        pca_aff_array[self_ind] = 1
        exp_weight = 0.5
        combined_aff_array = (1 - exp_weight)*spatial_aff_array + exp_weight*pca_aff_array
        col_arr = common_nn_cut
    else:
        pca_aff_value = pca_aff_mat_sym[i, pca_nn].toarray()[0]
        cut_off = np.quantile(pca_aff_value, (pca_aff_value.shape[0] - knn - 1)/pca_aff_value.shape[0])
        pca_nn_cut = pca_nn[pca_aff_value > cut_off]
        self_ind = np.where(pca_nn_cut == i)[0]
        pca_aff_array = pca_aff_mat_sym[i, pca_nn_cut].toarray()[0]
        pca_aff_array[np.where(pca_nn_cut != i)[0]] = pca_aff_array[np.where(pca_nn_cut != i)[0]]/np.sum(pca_aff_array[np.where(pca_nn_cut != i)[0]])
        pca_aff_array[self_ind] = 1
        combined_aff_array = pca_aff_array
        col_arr = pca_nn_cut
    row_arr = np.array([i]*combined_aff_array.shape[0])
    return(np.vstack([row_arr, col_arr, combined_aff_array]))

def read_h5(h5_file):
    expr_read = read_10X_h5(filename = h5_file)
    expr_mat = expr_read.matrix
    expr_genes = expr_read.names.tolist()
    expr_cells = expr_read.barcodes.tolist()
    if type(expr_genes[0]) == bytes:
        expr_genes = [i.decode() for i in expr_genes]
    if type(expr_cells[0]) == bytes:
        expr_cells = [i.decode() for i in expr_cells]
    return(expr_mat, expr_genes, expr_cells)

def Imputation(expr_file, spatial_file, num_workers, out_dir, out_prefix):
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    expr_mat_all, expr_genes_all, expr_cells_all = read_h5(expr_file)
    spatial_df = pd.read_csv(spatial_file, sep = "\t", header = 0, index_col = 0)
    spatial_df.index = spatial_df.index.astype(str)
    adata = ST_clustering(expr_mat_all, expr_genes_all, expr_cells_all, spatial_df, out_dir, out_prefix + "_original")
    imputed_mat_list = []
    cells_list = []
    adata_meta_df = adata.obs
    data_pca = adata.obsm['X_pca'][:,:20]
    data_pca = pd.DataFrame(data_pca, index = adata.obs.index)
    cell_types = adata_meta_df.loc[:, 'louvain_0.6'].unique()
    for cell_type in cell_types:
        print(cell_type)
        adata_meta_df_ct = adata_meta_df.loc[adata_meta_df.loc[:,'louvain_0.6'] == cell_type,:]
        cells_ct = adata_meta_df_ct.index
        expr_mat_ct, expr_genes_ct, expr_cells_ct = sub_mat(expr_mat_all, expr_genes_all, expr_cells_all, cells_sub = cells_ct.tolist())
        data_pca_ct = data_pca.loc[expr_cells_ct,:]
        data_spatial_ct = spatial_df.loc[expr_cells_ct, ['x', 'y']]
        n_samples = data_pca_ct.shape[0]
        n_trees = 50
        k_params = 50
        PCA_tree = BuildTree(data = data_pca_ct, n_trees = n_trees)
        res_list = []
        with concurrent.futures.ThreadPoolExecutor(max_workers = num_workers) as executor:
            for i in range(n_samples):
                arg_tuple = (PCA_tree, i, k_params)
                res_list.append(executor.submit(CalAffinity_PCA, *arg_tuple))
            done, not_done = concurrent.futures.wait(res_list, timeout=None)
            PCA_result_list = [future.result() for future in done]
        pca_aff_mat_sym = CalAffinityMat(PCA_result_list, n_samples)
        k_params = 50
        Spatial_tree = BuildTree(data = data_spatial_ct, n_trees = n_trees)
        res_list = []
        with concurrent.futures.ThreadPoolExecutor(max_workers = num_workers) as executor:
            for i in range(n_samples):
                arg_tuple = (Spatial_tree, i, k_params)
                res_list.append(executor.submit(CalAffinity_Spatial, *arg_tuple))
            done, not_done = concurrent.futures.wait(res_list, timeout=None)
            Spatial_result_list = [future.result() for future in done]
        spatial_aff_mat_sym = CalAffinityMat(Spatial_result_list, n_samples)
        knn = 10
        res_list = []
        with concurrent.futures.ThreadPoolExecutor(max_workers = num_workers) as executor:
            for i in range(n_samples):
                arg_tuple = (pca_aff_mat_sym, spatial_aff_mat_sym, knn, i)
                res_list.append(executor.submit(CombineAffinityMat, *arg_tuple))
            done, not_done = concurrent.futures.wait(res_list, timeout=None)
            Combine_result_list = [future.result() for future in done]
        Combine_aff_mat = CalCombineAffinityMat(Combine_result_list, n_samples)
        Combine_aff_mat = csc_matrix(Combine_aff_mat)
        expr_mat_ct_imputed = Combine_aff_mat.dot(expr_mat_ct.T)
        print(expr_mat_ct_imputed.shape)
        imputed_mat_list.append(expr_mat_ct_imputed.T)
        cells_list = cells_list + expr_cells_ct
    imputed_mat_merged = hstack(imputed_mat_list)
    expr_mat_imputed_all, expr_genes_imputed_all, expr_cells_imputed_all = sub_mat(imputed_mat_merged, expr_genes_ct, cells_list, genes_sub = expr_genes_all, cells_sub = expr_cells_all)
    out_combined_file = os.path.join(out_dir, "%s_k_params_%s_knn_%s_imputed.h5" %(out_prefix, k_params, knn))
    out_h5ad_file = os.path.join(out_dir, "%s_k_params_%s_knn_%s_imputed_anndata.h5ad" %(out_prefix, k_params, knn))
    write_10X_h5(out_combined_file, expr_mat_imputed_all, features = expr_genes_imputed_all, barcodes = expr_cells_imputed_all, datatype = 'Gene')
    adata_imputed = ST_clustering(expr_mat_imputed_all, expr_genes_imputed_all, expr_cells_imputed_all, spatial_df, out_dir, out_prefix + "_imputed")
    adata_imputed.write_h5ad(out_h5ad_file)

if __name__ == '__main__':
    parser = ImputationParser()
    expr_file = parser.expr
    spatial_file = parser.spatial
    num_workers = parser.num_workers
    out_dir = parser.out_dir
    out_prefix = parser.out_prefix
    Imputation(expr_file, spatial_file, num_workers, out_dir, out_prefix)
