# -*- coding: utf-8 -*-
# @Author: dongqing
# @Date:   2023-09-13 15:02:14
# @Last Modified by:   Dongqing
# @Last Modified time: 2025-08-26 14:02:59


# scale the distance, rather than the coordinates

import gc
import json
import os, sys
import argparse
import functools
import pandas as pd
import numpy as np
import multiprocessing as mp
import seaborn as sns
import skimage
from scipy.spatial import KDTree, distance
from sklearn.preprocessing import StandardScaler
import concurrent.futures
from random import sample, seed
from anndata import ImplicitModificationWarning

from Cellist.Utility import *
from Cellist.IO import *
from Cellist.Plot import *

import logging
logging.getLogger("fontTools.subset").setLevel(logging.WARNING)
logging.basicConfig(level=logging.INFO, format='%(asctime)s  %(message)s')

import warnings
pd.set_option('mode.chained_assignment', None)
warnings.filterwarnings('ignore', category=RuntimeWarning)
warnings.filterwarnings('ignore', category=ImplicitModificationWarning, append=True)

def CellistParser(subparsers):
    parser = subparsers.add_parser("seg", 
        help = "Run Cellist segmentation on high-resolution spatial transcriptomics. ")
    group_input = parser.add_argument_group("Input arguments")
    group_input.add_argument("--platform", dest = "platform", default = 'barcoding', choices = ['barcoding', 'imaging'],
        help = "Supported high-resoultion ST platform: 'barcoding' (e.g., Stereo-seq, Seq-scope) or "
        "'imaging' (e.g., seqFISH, STARmap, 10x Xenium). Default: barcoding.")
    group_input.add_argument("--resolution", dest = "resolution", type = float, default = 0.5,
        help = "Resolution of the platform (in microns), i.e., the distance between spatial units. "
        "For example, in Stereo-seq, where the distance between spots is 0.5 microns, the resolution should be set to 0.5. "
        "In seqFISH+, where the distance between pixels 103 nanometer, it should be set as 0.1. Default: 0.5. ")
    group_input.add_argument("--gem", dest = "spot_expr_file", default = None,
        help = "Bin1 gem file generated by SAW pipeline. ")
    group_input.add_argument("--spot-count-h5", dest = "all_spot_count_h5_file", default = None,
        help = "Bin1 spot-level expression file in h5 format.")
    group_input.add_argument("--nucleus-seg-method", dest = "nucleus_seg_method", default = "Watershed", choices = ['Watershed', 'Cellpose'],
        help = "The initial method to segment nuclei, Watershed or Cellpose. ")
    group_input.add_argument("--nucleus-prop", dest = "props_file", default = None,
        help = "Properties of segmented nucleus by Watershed or Cellpose algorithm. ")
    group_input.add_argument("--nucleus-count-h5", dest = "nucleus_count_h5_file", default = None,
        help = "Nuclei-level expression generated by Watershed or Cellpose segmentation with h5 format. ")
    group_input.add_argument("--nucleus-seg", dest = "nucleus_coord_file", default = None,
        help = "Spot-level Watershed or Cellpose segmentation result. ")
    group_input.add_argument("--nworkers", dest = "num_workers", type = int, default = 8,
        help = "Maximum number of workers to use. ")
    group_input.add_argument("--patch-data-dir", dest = "patch_data_dir", default = None,
        help = "Directory of patch data, for example, 'Data_HVG', generated by Cellist (not used in the first run of Cellist). ")

    group_model = parser.add_argument_group("Model arguments")
    group_model.add_argument("--cell-radius", dest = "max_dist", type = int, default = 15,
        help = "Expected cell radius (in microns). "
        "Considering the diffusion of transcripts, it should be 1.5X~2X larger than the real cell size. "
        "Default: 15")
    group_model.add_argument("--spot-imputation-distance", dest = "neigh_dist", type = float, default = 2.5,
        help = "Imputation distance for each spot (in microns). Default: 2.5. ")
    group_model.add_argument("--noise-prop", dest = "noise_prop", type = float, default = 0.25,
        help = "Proportion of noise spots. Default: 0.25. ")
    group_model.add_argument("--alpha", dest = "alpha", type = float, default = 0.8,
        help = "Weight for expression distance. Default: 0.8. ")
    group_model.add_argument("--sigma", dest = "sigma", type = float, default = 1.0,
        help = "Weight for regularization. Default: 1.0. ")
    group_model.add_argument("--beta", dest = "beta", type = float, default = 10,
        help = "Penalty weight. Default: 10. ")
    group_model.add_argument("--gene-use", dest = "gene_use", type = str, default = 'HVG', choices = ['HVG', 'Frequent'],
        help = "Gene to use to calculate spot-cell distance. "
        "'HVG' means highly variable genes, and 'Frequent' means top frequent genes. "
        "Default: HVG. ")
    group_model.add_argument("--two-step", dest = "two_step", action = "store_true",
        help = "Two-step mode. If set, Cellist will assign non-nuclei spots recusively. ")
    group_model.add_argument("--cyto", dest = "cyto", action = "store_true", 
        help = "Two-step cyto mode. If set, Cellist will calculate the correlation between non-nuclei spots and the 1st cycle of cytoplasm. ")
    group_model.add_argument("--cell-radius-s1-scale", dest = "max_dist_s1_scale",type = float,  default = 0.5,
        help = "Two-step mode. The allowed maximum distance scale in the first-step segmentation. ")
    group_model.add_argument("--noise-prop-s1", dest = "noise_prop_s1", type = float,  default = 0.4,
        help = "Proportion of noise spots in the first step. Default: 0.4. ")

    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 Update_centroid(phy_coord_df, cluster_col):
    phy_centroid = phy_coord_df.groupby([cluster_col]).mean()
    return(phy_centroid)

def Update_cell_expr(coord_df_slice, count_df_slice, count_name):
    count_df_slice = pd.merge(count_df_slice, coord_df_slice.loc[:, ('x_y', 'Cellist', 'Nucleus')], how = "inner", on = 'x_y')
    expr_mat, gene_list, cell_list = get_cell_mat(count_df_seg = count_df_slice, seg_res = 'Cellist', count_name = count_name)
    return(expr_mat, gene_list, cell_list)

def Assign_nonnucleispot(cell_expr_slice_array, phy_cell_centroid, nonnucl_enhanced_array_filtered, nonnucl_phy_coord_df,
    coord_df_slice, max_dist_slice, alpha, sigma, beta, noise_prop):
    nonnucl_spot_centroid_corr = pearson_corr(nonnucl_enhanced_array_filtered, cell_expr_slice_array)
    nonnucl_spot_centroid_corr[np.isnan(nonnucl_spot_centroid_corr)] = -1
    nonnucl_spot_centroid_dist_phy = distance.cdist(nonnucl_phy_coord_df, phy_cell_centroid, 'euclidean')
    # dist_scale_factor = np.max(nonnucl_spot_centroid_dist_phy)
    dist_scale_factor = 2*max_dist_slice
    P_dist = (nonnucl_spot_centroid_dist_phy - max_dist_slice)/dist_scale_factor
    P_dist[P_dist < 0] = 0
    R = np.exp(-(alpha*(1-nonnucl_spot_centroid_corr)/2 + (1-alpha)*nonnucl_spot_centroid_dist_phy/dist_scale_factor + beta*np.square(P_dist))/sigma)
    expr_loss = alpha*(1-nonnucl_spot_centroid_corr)/2
    dist_loss = (1-alpha)*nonnucl_spot_centroid_dist_phy/np.max(nonnucl_spot_centroid_dist_phy)
    regularization = sigma*np.log(R + 0.01)
    penalty = beta*np.square(P_dist)
    loss = expr_loss + dist_loss + regularization + penalty
    obj = np.nansum(np.multiply(R, loss))
    nonnucl_phy_coord_df.loc[:, 'Cellist'] = phy_cell_centroid.index[np.argmax(R, axis = 1)]
    nonnucl_phy_coord_df.loc[:, 'Cellist_prob'] = R.max(axis = 1)
    # nonnucl_phy_coord_df.loc[:, 'Cellist_slice'] = nonnucl_phy_coord_df.loc[:, 'Cellist'].copy()
    # q_noise = nonnucl_phy_coord_df['Cellist_prob'].quantile(noise_prop)
    # nonnucl_phy_coord_df.loc[nonnucl_phy_coord_df['Cellist_prob'] < q_noise, 'Cellist_slice'] = np.nan
    # nonnucl_spot_centroid_dist_phy_assigned = np.array([nonnucl_spot_centroid_dist_phy[i, np.argmax(R, axis = 1)[i]] for i in range(nonnucl_spot_centroid_dist_phy.shape[0])])
    # nonnucl_phy_coord_df.loc[nonnucl_spot_centroid_dist_phy_assigned > 2*max_dist_slice, 'Cellist_slice'] = np.nan
    R_df = pd.DataFrame(R, index = nonnucl_phy_coord_df.index, columns = phy_cell_centroid.index)
    R_sum = np.array(np.nansum(R, axis = 1).tolist()*R.shape[1]).reshape(R.shape, order = 'F')
    R_norm = np.divide(R, R_sum)
    R_norm[np.isnan(R_norm)] = 0
    nonnucl_phy_coord_df.loc[:, 'Cellist_prob_norm'] = R_norm.max(axis = 1)
    R_norm_df = pd.DataFrame(R_norm, index = nonnucl_phy_coord_df.index, columns = phy_cell_centroid.index)
    return(R_df, obj, nonnucl_phy_coord_df)

def Soft_assignment(nucleus_seg_method, nonnucl_enhanced_array, nonnucl_enhanced_genes, nonnucl_enhanced_spots, 
    coord_df_slice, count_df_slice, 
    max_dist_slice, alpha, sigma, beta, noise_prop):
    coord_df_slice.loc[:, 'Cellist'] = coord_df_slice[nucleus_seg_method].tolist()
    # not update cell centroid
    phy_cell_centroid = Update_centroid(phy_coord_df = coord_df_slice.loc[:, ('x', 'y', nucleus_seg_method)], cluster_col = nucleus_seg_method)
    count_name = count_df_slice.columns[3]
    try:
        cell_expr_slice_array, cell_expr_slice_gene, cell_expr_slice_cell = Update_cell_expr(coord_df_slice = coord_df_slice, count_df_slice = count_df_slice, count_name = count_name)
    except:
        return(None)
    gene_overlap = sorted(list(set(cell_expr_slice_gene) & set(nonnucl_enhanced_genes)))
    cell_expr_slice_array_filtered, cell_expr_slice_gene_filtered, cell_expr_slice_cell_filtered = sub_mat(
        mat = cell_expr_slice_array, genes = cell_expr_slice_gene, cells = cell_expr_slice_cell, 
        genes_sub = gene_overlap)
    nonnucl_enhanced_array_filtered, nonnucl_enhanced_genes_filtered, nonnucl_enhanced_spots_filtered = sub_mat(
        mat = nonnucl_enhanced_array, genes = nonnucl_enhanced_genes, cells = nonnucl_enhanced_spots, 
        genes_sub = cell_expr_slice_gene_filtered)
    coord_df_slice_nonnucl_df = coord_df_slice.loc[nonnucl_enhanced_spots_filtered, ('x', 'y')]
    phy_cell_centroid = phy_cell_centroid.loc[[float(i) for i in cell_expr_slice_cell_filtered],:]
    R_df, obj, nonnucl_phy_coord_df = Assign_nonnucleispot(
        cell_expr_slice_array = cell_expr_slice_array_filtered.toarray(),
        phy_cell_centroid = phy_cell_centroid, 
        nonnucl_enhanced_array_filtered = nonnucl_enhanced_array_filtered,
        nonnucl_phy_coord_df = coord_df_slice_nonnucl_df,
        coord_df_slice = coord_df_slice, max_dist_slice = max_dist_slice,
        alpha = alpha, sigma = sigma, beta = beta, noise_prop = noise_prop)
    coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist"] = nonnucl_phy_coord_df['Cellist'].tolist()
    coord_df_slice.loc[:, 'Cellist_prob'] = 1
    coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist_prob"] = nonnucl_phy_coord_df['Cellist_prob'].tolist()
    coord_df_slice.loc[coord_df_slice['Cellist_prob'] == 0.0, "Cellist"] = np.nan
    # coord_df_slice.loc[:, 'Cellist_slice'] = coord_df_slice.loc[:, 'Cellist'].copy()
    # coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist_slice"] = nonnucl_phy_coord_df['Cellist_slice'].tolist()
    coord_df_slice.loc[:, 'Cellist_prob_norm'] = 1
    coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist_prob_norm"] = nonnucl_phy_coord_df['Cellist_prob_norm'].tolist()
    coord_df_slice.loc[coord_df_slice['Cellist_prob_norm'] == 0.0, "Cellist_norm"] = np.nan
    return(coord_df_slice)

def Soft_assignment_2step(nucleus_seg_method, nonnucl_enhanced_array, nonnucl_enhanced_genes, nonnucl_enhanced_spots, 
    coord_df_slice, count_df_slice,
    max_dist_slice, alpha, sigma, beta, noise_prop, max_dist_s1_scale = 0.5, noise_prop_s1 = 0.4):
    coord_df_slice.loc[:, 'Cellist'] = coord_df_slice[nucleus_seg_method].tolist()
    # not update cell centroid
    multi_step = 2
    for i in range(multi_step):
        print('Segmentation %s.' %i)
        phy_cell_centroid = Update_centroid(phy_coord_df = coord_df_slice.loc[:, ('x', 'y', 'Cellist')], cluster_col = 'Cellist')
        count_name = count_df_slice.columns[3]
        try:
            cell_expr_slice_array, cell_expr_slice_gene, cell_expr_slice_cell = Update_cell_expr(coord_df_slice = coord_df_slice, count_df_slice = count_df_slice, count_name = count_name)
        except:
            return(None)
        gene_overlap = sorted(list(set(cell_expr_slice_gene) & set(nonnucl_enhanced_genes)))
        cell_expr_slice_array_filtered, cell_expr_slice_gene_filtered, cell_expr_slice_cell_filtered = sub_mat(
            mat = cell_expr_slice_array, genes = cell_expr_slice_gene, cells = cell_expr_slice_cell, 
            genes_sub = gene_overlap)
        coord_df_slice_nonnucl = coord_df_slice.loc[coord_df_slice['Cellist'].isna(),:]
        nonnucl_spots = sorted(list(set(nonnucl_enhanced_spots) & set(coord_df_slice_nonnucl.index.tolist())))
        nonnucl_enhanced_array_filtered, nonnucl_enhanced_genes_filtered, nonnucl_enhanced_spots_filtered = sub_mat(
            mat = nonnucl_enhanced_array, genes = nonnucl_enhanced_genes, cells = nonnucl_enhanced_spots, 
            genes_sub = cell_expr_slice_gene_filtered, cells_sub = nonnucl_spots)
        coord_df_slice_nonnucl_df = coord_df_slice.loc[nonnucl_enhanced_spots_filtered, ('x', 'y')]
        phy_cell_centroid = phy_cell_centroid.loc[[float(i) for i in cell_expr_slice_cell_filtered],:]
        if i == 0:
            max_dist_slice_step = max_dist_slice*max_dist_s1_scale
        else:
            max_dist_slice_step = max_dist_slice
        print(coord_df_slice_nonnucl_df.shape[0])
        R_df, obj, nonnucl_phy_coord_df = Assign_nonnucleispot(
            cell_expr_slice_array = cell_expr_slice_array_filtered.toarray(),
            phy_cell_centroid = phy_cell_centroid, 
            nonnucl_enhanced_array_filtered = nonnucl_enhanced_array_filtered,
            nonnucl_phy_coord_df = coord_df_slice_nonnucl_df,
            coord_df_slice = coord_df_slice, max_dist_slice = max_dist_slice_step,
            alpha = alpha, sigma = sigma, beta = beta, noise_prop = noise_prop)
        coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist"] = nonnucl_phy_coord_df['Cellist'].tolist()
        coord_df_slice.loc[:, 'Cellist_prob'] = 1
        coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist_prob"] = nonnucl_phy_coord_df['Cellist_prob'].tolist()
        coord_df_slice.loc[coord_df_slice['Cellist_prob'] == 0.0, "Cellist"] = np.nan
        coord_df_slice.loc[:, 'Cellist_prob_norm'] = 1
        coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist_prob_norm"] = nonnucl_phy_coord_df['Cellist_prob_norm'].tolist()
        coord_df_slice.loc[coord_df_slice['Cellist_prob_norm'] == 0.0, "Cellist_norm"] = np.nan
        if i == 0:
            coord_df_slice_cyto = coord_df_slice.loc[coord_df_slice['Nucleus'] == 0, ]
            q_noise = coord_df_slice_cyto['Cellist_prob'].quantile(noise_prop_s1)
            coord_df_slice.loc[coord_df_slice['Cellist_prob'] < q_noise, 'Cellist'] = np.nan
    return(coord_df_slice)

def Soft_assignment_2step_cyto(nucleus_seg_method, nonnucl_enhanced_array, nonnucl_enhanced_genes, nonnucl_enhanced_spots, 
    coord_df_slice_scaled_df, coord_df_slice, count_df_slice, 
    max_dist_slice_scaled, alpha, sigma, beta, prob_cutoff):
    coord_df_slice.loc[:, 'Cellist'] = coord_df_slice[nucleus_seg_method].tolist()
    multi_step = 2
    for i in range(multi_step):
        # cyto1-nucl
        print('Segmentation %s.' %i)
        coord_df_slice_scaled_df.loc[:,'Cellist'] = coord_df_slice.loc[:,'Cellist']
        phy_cell_centroid = Update_centroid(phy_coord_df = coord_df_slice_scaled_df.loc[:, ('x', 'y', 'Cellist')], cluster_col = 'Cellist')
        count_name = count_df_slice.columns[3]
        try:
            if i > 0:
                coord_df_slice_cyto1 = coord_df_slice.loc[(coord_df_slice['Cellist'].notna()) & (coord_df_slice['Nucleus'] == 0),:]
                cell_expr_slice_array, cell_expr_slice_gene, cell_expr_slice_cell = Update_cell_expr(coord_df_slice = coord_df_slice_cyto1, count_df_slice = count_df_slice, count_name = count_name)
            else:
                cell_expr_slice_array, cell_expr_slice_gene, cell_expr_slice_cell = Update_cell_expr(coord_df_slice = coord_df_slice, count_df_slice = count_df_slice, count_name = count_name)
        except:
            return(None)
        gene_overlap = sorted(list(set(cell_expr_slice_gene) & set(nonnucl_enhanced_genes)))
        cell_expr_slice_array_filtered, cell_expr_slice_gene_filtered, cell_expr_slice_cell_filtered = sub_mat(
            mat = cell_expr_slice_array, genes = cell_expr_slice_gene, cells = cell_expr_slice_cell, 
            genes_sub = gene_overlap)
        coord_df_slice_nonnucl = coord_df_slice.loc[coord_df_slice['Cellist'].isna(),:]
        nonnucl_spots = sorted(list(set(nonnucl_enhanced_spots) & set(coord_df_slice_nonnucl.index.tolist())))
        nonnucl_enhanced_array_filtered, nonnucl_enhanced_genes_filtered, nonnucl_enhanced_spots_filtered = sub_mat(
            mat = nonnucl_enhanced_array, genes = nonnucl_enhanced_genes, cells = nonnucl_enhanced_spots, 
            genes_sub = cell_expr_slice_gene_filtered, cells_sub = nonnucl_spots)
        coord_df_slice_nonnucl_scaled_df = coord_df_slice_scaled_df.loc[nonnucl_enhanced_spots_filtered, ('x', 'y')]
        phy_cell_centroid = phy_cell_centroid.loc[[float(i) for i in cell_expr_slice_cell_filtered],:]
        R_norm, obj, nonnucl_phy_coord_df = Assign_nonnucleispot(
            cell_expr_slice_array = cell_expr_slice_array_filtered.toarray(),
            phy_cell_centroid = phy_cell_centroid, 
            nonnucl_enhanced_array_filtered = nonnucl_enhanced_array_filtered,
            nonnucl_phy_coord_df = coord_df_slice_nonnucl_scaled_df,
            coord_df_slice = coord_df_slice, max_dist_slice_scaled = max_dist_slice_scaled/(multi_step-i),
            alpha = alpha, sigma = sigma, beta = beta, prob_cutoff = prob_cutoff)
        coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist"] = nonnucl_phy_coord_df['Cellist'].tolist()
        coord_df_slice.loc[:, 'Cellist_prob'] = 1
        coord_df_slice.loc[nonnucl_phy_coord_df.index.tolist(), "Cellist_prob"] = nonnucl_phy_coord_df['Cellist_prob'].tolist()
        # coord_df_slice.loc[coord_df_slice['Cellist_prob'].isna(), "Cellist"] = np.nan
        coord_df_slice.loc[coord_df_slice['Cellist_prob'] == 0.0, "Cellist"] = np.nan
    return(coord_df_slice)

def Soft_assignment_slice(count_df_patch, coord_df_patch, nucleus_seg_method, props_df_patch,
    all_count_mat, all_count_genes, all_count_spots, 
    xmin_cur, xmax_cur, ymin_cur, ymax_cur, selected_genes, max_dist = 15, resolution = 0.5,
    alpha = 0.8, out_dir = ".", sigma = 1.0, beta = 10, two_step = False, cyto = False, max_dist_s1_scale = 0.5, noise_prop_s1 = 0.4, noise_prop = 0.25, neigh_dist = 2.5):
    coord_df_slice = coord_df_patch.loc[(coord_df_patch['x'] >= xmin_cur) & (coord_df_patch['x'] < xmax_cur) &
                                (coord_df_patch['y'] >= ymin_cur) & (coord_df_patch['y'] < ymax_cur) , :]
    count_df_slice = count_df_patch.loc[(count_df_patch['x'] >= xmin_cur) & (count_df_patch['x'] < xmax_cur) &
                                (count_df_patch['y'] >= ymin_cur) & (count_df_patch['y'] < ymax_cur) , :]
    coord_df_slice_nonnucl = coord_df_slice[coord_df_slice['Nucleus'] == 0]
    coord_df_slice_nucl = coord_df_slice[coord_df_slice['Nucleus'] == 1]
    nucl_spot_count = coord_df_slice_nucl[nucleus_seg_method].value_counts()
    props_df_slice = props_df_patch.loc[(props_df_patch['centroid-0'] >= xmin_cur) & (props_df_patch['centroid-0'] < xmax_cur) & 
                                (props_df_patch['centroid-1'] >= ymin_cur) & (props_df_patch['centroid-1'] < ymax_cur), :]
    if coord_df_slice_nucl.shape[0] > 0 and props_df_slice.shape[0] > 0:
        props_df_slice.loc[np.intersect1d(props_df_slice.index, nucl_spot_count.index),"nSpot"] = nucl_spot_count[np.intersect1d(props_df_slice.index, nucl_spot_count.index)]
        props_df_slice = props_df_slice.loc[props_df_slice['nSpot'] >= 20, :]
        if props_df_slice.shape[0] > 0:
            nucl_cell_slice = list(set(coord_df_slice_nucl[nucleus_seg_method].unique()) & set(props_df_slice.index.tolist()))
            nucl_cell_slice = sorted(nucl_cell_slice)
            nucl_cell_rm = set(coord_df_slice_nucl[nucleus_seg_method].unique()) - set(nucl_cell_slice)
            if nucl_cell_slice:
                coord_df_slice.loc[coord_df_slice[nucleus_seg_method].isin(nucl_cell_rm), nucleus_seg_method] = np.nan
                spot_expr_slice_filter, spot_gene_slice_filter, spot_slice_filter = sub_mat(
                    mat = all_count_mat, genes = all_count_genes, cells = all_count_spots, 
                    cells_sub = coord_df_slice.index.tolist(), 
                    genes_sub = selected_genes)
                nonnucl_spots = sorted(list(set(spot_slice_filter) & set(coord_df_slice_nonnucl.index.tolist())))
                if nonnucl_spots:
                    # enhance single spot's expr
                    neigh_dist_slice = int(neigh_dist/resolution)
                    spot_expr_slice_filter_enhanced_array = enhance_spot_expr(coord_df = coord_df_slice.loc[spot_slice_filter, :], 
                        coord_expr_mat = spot_expr_slice_filter, neigh_dist = neigh_dist_slice)
                    spot_expr_slice_filter_enhanced_array_nonnucl, spot_expr_slice_filter_enhanced_array_nonnucl_gene, spot_expr_slice_filter_enhanced_array_nonnucl_spot = sub_mat(
                        mat = spot_expr_slice_filter_enhanced_array,
                        genes = spot_gene_slice_filter, 
                        cells = spot_slice_filter,
                        cells_sub = nonnucl_spots)
                    coord_df_slice_nonnucl = coord_df_slice_nonnucl.loc[spot_expr_slice_filter_enhanced_array_nonnucl_spot, :]
                    if type(max_dist) == int:
                        max_dist_slice = max_dist/resolution
                    else: # nuclei size
                        max_dist_slice = 1.5*props_df_slice['equivalent_diameter_area'].median()
                    if two_step:
                        if not cyto:
                            coord_df_slice = Soft_assignment_2step(
                                nucleus_seg_method = nucleus_seg_method,
                                nonnucl_enhanced_array = spot_expr_slice_filter_enhanced_array_nonnucl,
                                nonnucl_enhanced_genes = spot_expr_slice_filter_enhanced_array_nonnucl_gene,
                                nonnucl_enhanced_spots = spot_expr_slice_filter_enhanced_array_nonnucl_spot,
                                coord_df_slice = coord_df_slice,
                                count_df_slice = count_df_slice, 
                                max_dist_slice = max_dist_slice,
                                alpha = alpha, sigma = sigma, beta = beta, 
                                noise_prop = noise_prop,
                                max_dist_s1_scale = max_dist_s1_scale, 
                                noise_prop_s1 = noise_prop_s1)
                        else:
                            coord_df_slice = Soft_assignment_2step_cyto(
                                nucleus_seg_method = nucleus_seg_method,
                                nonnucl_enhanced_array = spot_expr_slice_filter_enhanced_array_nonnucl,
                                nonnucl_enhanced_genes = spot_expr_slice_filter_enhanced_array_nonnucl_gene,
                                nonnucl_enhanced_spots = spot_expr_slice_filter_enhanced_array_nonnucl_spot,
                                coord_df_slice_scaled_df = coord_df_slice_scaled_df, 
                                coord_df_slice = coord_df_slice,
                                count_df_slice = count_df_slice, 
                                max_dist_slice_scaled = max_dist_slice_scaled,
                                alpha = alpha, sigma = sigma, beta = beta, 
                                prob_cutoff = prob_cutoff)
                    else:
                        coord_df_slice = Soft_assignment(
                            nucleus_seg_method = nucleus_seg_method,
                            nonnucl_enhanced_array = spot_expr_slice_filter_enhanced_array_nonnucl,
                            nonnucl_enhanced_genes = spot_expr_slice_filter_enhanced_array_nonnucl_gene,
                            nonnucl_enhanced_spots = spot_expr_slice_filter_enhanced_array_nonnucl_spot,
                            coord_df_slice = coord_df_slice,
                            count_df_slice = count_df_slice, 
                            max_dist_slice = max_dist_slice,
                            alpha = alpha, sigma = sigma, beta = beta, noise_prop = noise_prop)
                    if type(coord_df_slice) == pd.DataFrame:
                        tmp_dir = os.path.join(out_dir, "tmp")
                        out_file_tmp = os.path.join(tmp_dir, "Sub_%s_%s_%s_%s_soft_clustering.txt" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur))
                        coord_df_slice.to_csv(out_file_tmp, index = False, sep = "\t")
                    else:
                        print("%s:%s:%s:%s:empty" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur))
                        out_file_tmp = "NULL"
                else:
                    coord_df_slice['Cellist'] = coord_df_slice[nucleus_seg_method]
                    coord_df_slice['Cellist_prob'] = 1
                    tmp_dir = os.path.join(out_dir, "tmp")
                    out_file_tmp = os.path.join(tmp_dir, "Sub_%s_%s_%s_%s_soft_clustering.txt" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur))
                    coord_df_slice.to_csv(out_file_tmp, index = False, sep = "\t")
            else:
                print("%s:%s:%s:%s:empty" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur))
                out_file_tmp = "NULL"
        else:
            print("%s:%s:%s:%s:empty" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur))
            out_file_tmp = "NULL"
    else:
        print("%s:%s:%s:%s:empty" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur))
        out_file_tmp = "NULL"
    return(out_file_tmp)


def Write_patch(count_df, coord_df, nucleus_seg_method, props_df, nucleus_expr_mat, nucleus_gene, nucleus_cell, all_count_mat, all_count_genes, all_count_spots,
    xmin_cur, xmax_cur, ymin_cur, ymax_cur,
    patch_data_dir, gene_use = 'HVG'):
    # get small patches to set parameters
    coord_df_patch = coord_df.loc[(coord_df['x'] >= xmin_cur) & (coord_df['x'] < xmax_cur) &
                                (coord_df['y'] >= ymin_cur) & (coord_df['y'] < ymax_cur) , :]
    coord_df_patch_nucl = coord_df_patch[coord_df_patch['Nucleus'] == 1]
    props_df_patch = props_df.loc[(props_df['centroid-0'] >= xmin_cur) & (props_df['centroid-0'] < xmax_cur) & 
                                (props_df['centroid-1'] >= ymin_cur) & (props_df['centroid-1'] < ymax_cur), :]
    count_df_patch = count_df.loc[(count_df['x'] >= xmin_cur) & (count_df['x'] < xmax_cur) &
                                (count_df['y'] >= ymin_cur) & (count_df['y'] < ymax_cur) , :]
    patch_prefix = "Patch_%s_%s_%s_%s" %(xmin_cur, xmax_cur, ymin_cur, ymax_cur)
    if coord_df_patch_nucl.shape[0] > 0 and props_df_patch.shape[0] > 0:
        nucl_cell_patch = list(set(coord_df_patch_nucl[nucleus_seg_method].unique()) & set(props_df_patch.index.tolist()))
        nucl_cell_patch = sorted(nucl_cell_patch)
        logging.info('There are %s cells in this patch.' %(len(nucl_cell_patch)))
        if len(nucl_cell_patch) > 0:
            nucleus_expr_sub, nucleus_gene_sub, nucleus_cell_sub = sub_mat(mat = nucleus_expr_mat, genes = nucleus_gene, cells = nucleus_cell, 
                cells_sub = nucl_cell_patch)
            if gene_use == "HVG":
                # get highly varibale genes
                try:
                    selected_genes = get_hvg(nucleus_expr_sub, nucleus_gene_sub, nucleus_cell_sub, 1500)
                except:
                    logging.info('Fail to identify HVG in %s.' %patch_prefix)
                    return('NULL')
            else:
                selected_genes = get_frequent_gene(nucleus_expr_sub, nucleus_gene_sub, nucleus_cell_sub, 0.05)
            del nucleus_expr_sub, nucleus_gene_sub, nucleus_cell_sub
            gc.collect()
            try:
                spot_expr_patch_filter, spot_gene_patch_filter, spot_patch = sub_mat(
                        mat = all_count_mat, genes = all_count_genes, cells = all_count_spots, 
                        cells_sub = coord_df_patch.index.tolist(), 
                        genes_sub = selected_genes)
            except:
                sys.exit('Fail to subset matrix in %s.' %patch_prefix)
            Write_tmp_files(count_df_patch, coord_df_patch, props_df_patch, spot_expr_patch_filter, spot_gene_patch_filter, spot_patch, patch_prefix, patch_data_dir)
            return(patch_prefix)
        else:
            return('NULL')
    else:
        return('NULL')

def Soft_assignment_patch(platform, resolution, nucleus_seg_method, props_file_patch, coord_file_patch, all_spot_count_h5_file_patch, count_file_patch,
    out_dir, num_workers, 
    alpha, sigma, beta, gene_use, max_dist, two_step, cyto, max_dist_s1_scale, noise_prop_s1, noise_prop, neigh_dist):
    props_df_patch = pd.read_csv(props_file_patch, sep = "\t", header = 0, index_col = 0)
    coord_df_patch = pd.read_csv(coord_file_patch, sep = "\t", header = 0, index_col = 0)
    count_df_patch = pd.read_csv(count_file_patch, sep = "\t", header = 0)
    spot_expr_patch_filter, spot_gene_patch_filter, spot_patch = read_h5(all_spot_count_h5_file_patch)
    patch_prefix = os.path.basename(props_file_patch).split("Patch_")[1].split('_property.txt')[0]
    if platform == "barcoding":
        slice_size = 200
        stride_length = 10
    elif platform == "imaging":
        slice_size = 2200
        stride_length = 100
    selected_genes = spot_gene_patch_filter
    coord_slice_list = slice_coord(coord_df_patch, slice_size, stride_length)
    logging.info('There are %s slices in the patch.' %len(coord_slice_list))
    res_list = []
    with concurrent.futures.ThreadPoolExecutor(max_workers = 36) as executor:
        for coord_tuple in coord_slice_list:
            xmin_slice, xmax_slice, ymin_slice, ymax_slice = coord_tuple
            arg_tuple = (count_df_patch, coord_df_patch, nucleus_seg_method, props_df_patch, spot_expr_patch_filter, 
                spot_gene_patch_filter, spot_patch, xmin_slice, xmax_slice, ymin_slice, ymax_slice, 
                selected_genes, max_dist, resolution, alpha, out_dir, sigma, beta, two_step, cyto,max_dist_s1_scale, noise_prop_s1, noise_prop, neigh_dist)
            res_list.append(executor.submit(Soft_assignment_slice, *arg_tuple))
        done, not_done = concurrent.futures.wait(res_list, timeout=None)
        tmp_file_list = [future.result() for future in done]
    return(tmp_file_list)


def Read_files(nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, all_spot_count_h5_file, spot_expr_file):
    # ---------- Read files ------------
    # Read property file
    props_df = pd.read_csv(props_file, sep = "\t")
    props_df.index = props_df['label'].astype(np.float64)
    props_df.index.name = None
    # Read nucleus expr file
    nucleus_expr_mat, nucleus_gene, nucleus_cell = read_h5(nucleus_count_h5_file)
    nucleus_cell = [float(i) for i in nucleus_cell]
    # Read nucleus file
    coord_df = pd.read_csv(nucleus_coord_file, sep = "\t", header = 0)
    coord_df.index = coord_df['x_y']
    coord_df.index.name = None
    coord_df.loc[coord_df['Nucleus'] == 0, nucleus_seg_method] = np.nan
    # Read spot-level expr file (with nucleus 0/1)]
    count_df = pd.read_csv(spot_expr_file, sep = "\t", comment = "#")
    count_df.loc[:, 'x_y'] = count_df['x'].astype(str) + '_' + count_df['y'].astype(str)
    # Read spot-level expr mat
    all_count_mat, all_count_genes, all_count_spots = read_h5(all_spot_count_h5_file)
    return(count_df, coord_df, props_df, 
        nucleus_expr_mat, nucleus_gene, nucleus_cell, 
        all_count_mat, all_count_genes, all_count_spots)

def Write_tmp_files(count_df_patch, coord_df_patch, props_df_patch, spot_expr_patch_filter, spot_gene_patch_filter, spot_patch, 
    patch_prefix, patch_data_dir):
    # ---------- Write tmp patch files ------------
    # Write property file
    props_df_patch.to_csv(os.path.join(patch_data_dir, "%s_property.txt" %patch_prefix), sep = "\t")
    # Write spot coord file
    coord_df_patch.to_csv(os.path.join(patch_data_dir, "%s_coord.txt" %patch_prefix), sep = "\t")
    # Write spot-level expr file (with nucleus 0/1)]
    count_df_patch.to_csv(os.path.join(patch_data_dir, "%s_count.txt" %patch_prefix), sep = "\t", index = False)
    # Write spot-level expr mat
    write_10X_h5(os.path.join(patch_data_dir, "%s_spot_count.h5" %patch_prefix), spot_expr_patch_filter, spot_gene_patch_filter, spot_patch, datatype = 'Gene')

def Assign_filter(coord_df, out_dir, out_prefix, noise_prop = 0.25):
    coord_df_nucl = coord_df.loc[coord_df['Nucleus']==1, ]
    coord_df_cyto = coord_df.loc[coord_df['Nucleus']==0, ]
    # hist_res = plt.hist(coord_df_cyto['Cellist_prob'], bins = 200)
    # hist_center = (hist_res[1][:-1] + hist_res[1][1:])/2
    # otsu2_thre = skimage.filters.threshold_otsu(hist = (hist_res[0], hist_center))
    # otsu3_thre = skimage.filters.threshold_multiotsu(classes=3, hist = (hist_res[0], hist_center))
    f, ax = plt.subplots(1,1, figsize=(4, 1.6))
    line1 = sns.kdeplot(data = coord_df_cyto, 
        x = "Cellist_prob", ax = ax, color = "#C00000")
    line2= ax.axvline(x = coord_df_cyto['Cellist_prob'].quantile(0.25), color="black", linestyle="--", linewidth = 0.9)
    # line3= ax.axvline(x = otsu2_thre, color="#009292", linestyle=":", linewidth = 1)
    # line4= ax.axvline(x = otsu3_thre[0], color="#FFB6DB", linestyle="-", linewidth = 1)
    # line5= ax.axvline(x = otsu3_thre[1], color="#FF6DB6", linestyle="-", linewidth = 1)
    # plt.legend(handles = [line2, line3, line4, line5], labels = ['Q1', 'Otsu threshold (binary)', '3-class, lower', '3-class, upper'],
    #     fontsize = 7.5, markerscale = 0.8)
    plt.xlabel('Assign probability')
    plt.title(out_prefix, pad = 10)
    plot_file = os.path.join(out_dir, "%s_assign_probability.pdf" %out_prefix)
    plt.savefig(plot_file, bbox_inches='tight')
    # norm
    # hist_res = plt.hist(coord_df_cyto['Cellist_prob_norm'], bins = 200)
    # hist_center = (hist_res[1][:-1] + hist_res[1][1:])/2
    # otsu2_thre = skimage.filters.threshold_otsu(hist = (hist_res[0], hist_center))
    # otsu3_thre = skimage.filters.threshold_multiotsu(classes=3, hist = (hist_res[0], hist_center))
    f, ax = plt.subplots(1,1, figsize=(4, 1.6))
    line1 = sns.kdeplot(data = coord_df_cyto, 
        x = "Cellist_prob_norm", ax = ax, color = "#C00000")
    line2= ax.axvline(x = coord_df_cyto['Cellist_prob_norm'].quantile(0.25), color="black", linestyle="--", linewidth = 0.9)
    # line3= ax.axvline(x = otsu2_thre, color="#009292", linestyle=":", linewidth = 1)
    # line4= ax.axvline(x = otsu3_thre[0], color="#FFB6DB", linestyle="-", linewidth = 1)
    # line5= ax.axvline(x = otsu3_thre[1], color="#FF6DB6", linestyle="-", linewidth = 1)
    # plt.legend(handles = [line2, line3, line4, line5], labels = ['Q1', 'Otsu threshold (binary)', '3-class, lower', '3-class, upper'],
    #     fontsize = 7.5, markerscale = 0.8)
    plt.xlabel('Assign probability')
    plt.title(out_prefix, pad = 10)
    plot_file = os.path.join(out_dir, "%s_assign_probability_normalized.pdf" %out_prefix)
    plt.savefig(plot_file, bbox_inches='tight')
    q_noise = coord_df_cyto['Cellist_prob'].quantile(noise_prop)
    coord_df['Cellist_unfiltered'] = coord_df['Cellist'].copy()
    coord_df.loc[coord_df['Cellist_prob'] < q_noise, 'Cellist'] = np.nan
    # coord_df.loc[(coord_df['Cellist'].isna()) & (coord_df['Cellist_slice'].notna()), 'Cellist'] = coord_df.loc[(coord_df['Cellist'].isna()) & (coord_df['Cellist_slice'].notna()), 'Cellist_slice']
    q_noise_norm = coord_df_cyto['Cellist_prob_norm'].quantile(noise_prop)
    coord_df['Cellist_norm'] = coord_df['Cellist_unfiltered'].copy()
    coord_df.loc[coord_df['Cellist_prob_norm'] < q_noise_norm, 'Cellist_norm'] = np.nan
    return(coord_df)

def Cellist(platform, resolution, nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, all_spot_count_h5_file, spot_expr_file,
    patch_data_dir, num_workers,
    alpha, sigma, beta, gene_use, max_dist, two_step, cyto, max_dist_s1_scale, noise_prop_s1, noise_prop, neigh_dist, out_dir, out_prefix):
    # ----------- output prefix ----------
    out_prefix = out_prefix + "_Cellist"
    if platform == 'barcoding':
        patch_size = 1100
        stride_length = 100
    elif platform == 'imaging':
        patch_size = 2200
        stride_length = 100
    if two_step:
        out_para = "alpha_%s_sigma_%s_beta_%s_gene_%s_dist_%s_twostep_%s_cyto_%s_noise_%s_neigh_%s_dscale_%s_s1noise_%s" %(alpha, sigma, beta, gene_use, max_dist, two_step, cyto, noise_prop, neigh_dist, max_dist_s1_scale, noise_prop_s1)
    else:
        out_para = "alpha_%s_sigma_%s_beta_%s_gene_%s_dist_%s_twostep_%s_cyto_%s_noise_%s_neigh_%s" %(alpha, sigma, beta, gene_use, max_dist, two_step, cyto, noise_prop, neigh_dist)
    out_dir = os.path.join(out_dir, out_para)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    tmp_dir = os.path.join(out_dir, "tmp")
    if not os.path.exists(tmp_dir):
        os.makedirs(tmp_dir)
    # -------------- write tmp patch files -------------
    if patch_data_dir:
        patch_prefix_list = []
        for patch_file in os.listdir(patch_data_dir):
            if patch_file.endswith('_spot_count.h5'):
                patch_prefix = patch_file.split('_spot_count.h5')[0]
                patch_prefix_list.append(patch_prefix)
    else:
        logging.info('Reading input files...')
        count_df, coord_df, props_df, nucleus_expr_mat, nucleus_gene, nucleus_cell, all_count_mat, all_count_genes, all_count_spots = Read_files(
                nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, all_spot_count_h5_file, spot_expr_file)
        # # ----------- slice patches ----------
        coord_patch_list = slice_coord(coord_df, patch_size, stride_length)
        logging.info('There are %s patches to be processed!' %len(coord_patch_list))
        patch_data_dir = os.path.join(out_dir, "..", "Data_%s" %gene_use)
        if not os.path.exists(patch_data_dir):
            os.mkdir(patch_data_dir)
        res_list = []
        with concurrent.futures.ThreadPoolExecutor(max_workers = 36) as executor:
            for i in range(len(coord_patch_list)):
                logging.info('Writing patch %s...' %i)
                xmin_patch, xmax_patch, ymin_patch, ymax_patch = coord_patch_list[i]
                arg_tuple = (count_df, coord_df, nucleus_seg_method, props_df, nucleus_expr_mat, nucleus_gene, nucleus_cell, all_count_mat, all_count_genes, all_count_spots,
                xmin_patch, xmax_patch, ymin_patch, ymax_patch,
                patch_data_dir, gene_use)
                res_list.append(executor.submit(Write_patch, *arg_tuple))
            done, not_done = concurrent.futures.wait(res_list, timeout=None)
            tmp_input_list_all = [future.result() for future in done]
        del count_df, coord_df, props_df, nucleus_expr_mat, nucleus_gene, nucleus_cell, all_count_mat, all_count_genes, all_count_spots
        gc.collect()
        patch_prefix_list = []
        for patch in tmp_input_list_all:
            if patch != 'NULL':
                patch_prefix_list.append(patch)
    # ----------- segmentation ----------
    logging.info('Start segmentation...')
    patch_res_list = []
    tmp_file_list_all = []
    with concurrent.futures.ProcessPoolExecutor(max_workers = num_workers) as executor_patch:
        for i in range(len(patch_prefix_list)):
            patch_prefix = patch_prefix_list[i]
            logging.info('Processing patch %s...' %patch_prefix)
            props_file_patch = os.path.join(patch_data_dir, "%s_property.txt" %patch_prefix)
            coord_file_patch = os.path.join(patch_data_dir, "%s_coord.txt" %patch_prefix)
            all_spot_count_h5_file_patch = os.path.join(patch_data_dir, "%s_spot_count.h5" %patch_prefix)
            count_file_patch = os.path.join(patch_data_dir, "%s_count.txt" %patch_prefix)
            arg_tuple_patch = (platform, resolution, nucleus_seg_method, props_file_patch, coord_file_patch, all_spot_count_h5_file_patch, count_file_patch,
                out_dir, num_workers, 
                alpha, sigma, beta, gene_use, max_dist, two_step, cyto, max_dist_s1_scale, noise_prop_s1, noise_prop, neigh_dist)
            patch_res_list.append(executor_patch.submit(Soft_assignment_patch, *arg_tuple_patch))
        done_patch, not_done_patch = concurrent.futures.wait(patch_res_list, timeout=None)
        for future in done_patch:
            tmp_file_list_all = tmp_file_list_all + future.result()
    # ------------ merge results from patches ------------
    logging.info('Merging results from patches...')
    count_df, coord_df, _, nucleus_expr_mat, nucleus_gene, nucleus_cell, _, _, _ = Read_files(
            nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, all_spot_count_h5_file, spot_expr_file)
    coord_df_sub_list = []
    for tmp_file in tmp_file_list_all:
        if tmp_file != "NULL":
            coord_df_sub = pd.read_csv(tmp_file, sep = "\t")
            coord_df_sub_list.append(coord_df_sub)
    coord_df_merge = pd.concat(coord_df_sub_list)
    coord_df_merge = coord_df_merge.drop_duplicates(subset = ['x_y', 'Cellist'])
    coord_df_merge = coord_df_merge.dropna(subset = ['Cellist'])
    coord_df_merge = coord_df_merge.sort_values(by=['x_y', 'Cellist_prob'])
    coord_df_merge = coord_df_merge.drop_duplicates(subset = ['x_y'], keep = 'last')
    # coord_df_merge = pd.merge(coord_df, coord_df_merge[['x_y', 'Cellist', 'Cellist_prob', 'Cellist_norm', 'Cellist_prob_norm', 'Cellist_slice']], how = "left", on = "x_y")
    coord_df_merge = pd.merge(coord_df, coord_df_merge[['x_y', 'Cellist', 'Cellist_prob', 'Cellist_norm', 'Cellist_prob_norm']], how = "left", on = "x_y")
    coord_df_merge = Assign_filter(coord_df_merge, out_dir, out_prefix, noise_prop)
    logging.info('Writing segmentation results...')
    coord_df_merge.to_csv(os.path.join(out_dir, "%s_segmentation.txt" %(out_prefix)), sep = "\t", index = False)
    draw_segmentation(coord_df_sub = coord_df_merge, seg_res = "Cellist", out_prefix = out_prefix, 
    out_dir = out_dir, x = "x", y = "y", figsize = (80, 80))
    count_df = pd.merge(count_df, coord_df_merge[['x_y', 'Cellist', 'Nucleus']], how = "left", on = "x_y")
    count_name = count_df.columns[3]
    write_segmentation_h5(count_df, "Cellist", out_prefix, out_dir, count_name)
    write_segmentation_cell_coord(coord_df_merge, "Cellist", out_prefix, out_dir)
    logging.info('Writing parameters and basic cell statistics...')
    hvg_list = get_hvg(nucleus_expr_mat, nucleus_gene, nucleus_cell, 1000)
    cell_num = coord_df_merge['Cellist'].value_counts().shape[0]
    nspot_all = int(coord_df_merge['Cellist'].count())
    nspot_nucleus = int(coord_df_merge[nucleus_seg_method].count())
    nspot_avg = int(nspot_all/cell_num)
    para_dict = {
        'nucleus_seg_method': nucleus_seg_method,
        'props_file': props_file, 
        'nucleus_count_h5_file': nucleus_count_h5_file, 
        'nucleus_coord_file': nucleus_coord_file, 
        'all_spot_count_h5_file': all_spot_count_h5_file, 
        'spot_expr_file': spot_expr_file, 
        'num_workers': num_workers, 
        'alpha': alpha, 
        'sigma': sigma, 
        'beta': beta, 
        'gene_use': gene_use, 
        'max_dist': max_dist, 
        'two_step': two_step,
        'noise_prop': noise_prop,
        'cell_num': cell_num,
        'nspot_nucleus': nspot_nucleus,
        'nspot_all': nspot_all,
        'nspot_avg': nspot_avg
    }
    try:
        if len(nucleus_cell) > 100000:
            seed(0)
            cells_sample = sample(nucleus_cell, 100000)
            Cellist_corr_nucl_cyto, Cellist_corr_nucl_cyto_cell = cal_corr_within_seg_nucleus(count_df.loc[count_df['Cellist'].isin(cells_sample), :], "Cellist", hvg_list, count_name)
        else:
            Cellist_corr_nucl_cyto, Cellist_corr_nucl_cyto_cell = cal_corr_within_seg_nucleus(count_df, "Cellist", hvg_list, count_name)
        Cellist_corr_nucl_cyto_df = pd.DataFrame({'Cell': Cellist_corr_nucl_cyto_cell, 'Correlation': Cellist_corr_nucl_cyto})
        Cellist_corr_nucl_cyto_df.to_csv(os.path.join(out_dir, "%s_corr_nucl_cyto_df.txt" %out_prefix), index = False, sep = "\t")
        # statics
        para_dict['Correlation'] = np.nanquantile(Cellist_corr_nucl_cyto, [0, 0.25, 0.5, 0.75, 1]).tolist()
    except:
        pass
        # write parameters
    with open(os.path.join(out_dir, "parameters.json"), "w") as outfile:
        json.dump(para_dict, outfile)

if __name__ == '__main__':
    parser = CellistParser()
    platform = parser.platform
    resolution = parser.resolution
    spot_expr_file = parser.spot_expr_file
    all_spot_count_h5_file = parser.all_spot_count_h5_file
    nucleus_seg_method = parser.nucleus_seg_method
    props_file = parser.props_file
    nucleus_count_h5_file = parser.nucleus_count_h5_file
    nucleus_coord_file = parser.nucleus_coord_file
    patch_data_dir = parser.patch_data_dir
    num_workers = parser.num_workers
    max_dist = parser.max_dist
    neigh_dist = parser.neigh_dist
    noise_prop = parser.noise_prop
    alpha = parser.alpha
    sigma = parser.sigma
    beta = parser.beta
    gene_use = parser.gene_use
    two_step = parser.two_step
    cyto = parser.cyto
    out_dir = parser.out_dir
    out_prefix = parser.out_prefix
    max_dist_s1_scale = parser.max_dist_s1_scale
    noise_prop_s1 = parser.noise_prop_s1
    Cellist(platform, resolution, nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, all_spot_count_h5_file, spot_expr_file,
        patch_data_dir, num_workers, alpha, sigma, beta, gene_use, max_dist, two_step, cyto, max_dist_s1_scale, noise_prop_s1, noise_prop, neigh_dist,
        out_dir, out_prefix)
