# -*- coding: utf-8 -*-
# @Author: dongqing
# @Date:   2025-04-17 13:26:55
# @Last Modified by:   dongqing
# @Last Modified time: 2025-04-17 14:15:49


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 *
from Cellist.Segmentation import Read_files, Assign_filter

import logging
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 TuneParser():
    parser = argparse.ArgumentParser(
                        description='Tune the --noise_prop based on the segmentation results with other parameters fixed. ')
    group_input = parser.add_argument_group("Input arguments")
    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("--cellist-seg", dest = "cellist_coord_file", default = None,
        help = "Existing spot-level Cellist segmentation result. ")

    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_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 Tune(nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, cellist_coord_file, all_spot_count_h5_file, spot_expr_file,
    alpha, sigma, beta, gene_use, max_dist, two_step, cyto, noise_prop, neigh_dist, out_dir, out_prefix):
    # ----------- output prefix ----------
    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)
    # ------------ 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_merge = pd.read_csv(cellist_coord_file , sep = "\t")
    coord_df_merge['Cellist'] = coord_df_merge['Cellist_unfiltered'].copy()
    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, 
        '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_cellist_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 = TuneParser().parse_args()
    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
    cellist_coord_file = parser.cellist_coord_file
    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
    Tune(nucleus_seg_method, props_file, nucleus_count_h5_file, nucleus_coord_file, cellist_coord_file, all_spot_count_h5_file, spot_expr_file,
        alpha, sigma, beta, gene_use, max_dist, two_step, cyto, noise_prop, neigh_dist,
        out_dir, out_prefix)

