# -*- coding: utf-8 -*-
# @Author: dongqing
# @Date:   2025-03-08 16:01:35
# @Last Modified by:   Dongqing
# @Last Modified time: 2025-08-26 13:15:17


import sys, os, gc
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

import torch
import matplotlib
matplotlib.rcParams['font.family'] = 'sans-serif'
matplotlib.rcParams['font.sans-serif'] = ['Arial']

import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s  %(message)s')

from cellpose import core, utils, io, models, metrics

from skimage.io import imread, imsave
from skimage.filters import threshold_otsu, threshold_multiotsu, threshold_niblack, threshold_local
from skimage.feature import peak_local_max
from skimage.measure import label, regionprops_table
from skimage.segmentation import watershed, expand_labels, mark_boundaries
from skimage.color import label2rgb, rgb2gray, rgb2hed
from skimage.segmentation import mark_boundaries
from skimage.morphology import remove_small_holes

from scipy.sparse import csc_matrix, save_npz, load_npz, csr_matrix
from scipy import ndimage as ndi

from Cellist.Utility import *
from Cellist.IO import gem_to_mat
from Cellist.Plot import *
from Cellist.Watershed import Thresholding, write_segmentation_coord, write_segmentation_cell

use_GPU = core.use_gpu()
print('>>> GPU activated? %s'%use_GPU)

'''
cellist watershed --gem /mnt/Venus/home/dongxin/Projects/heart/202401_Stereo_seq/result_C02942B3/041.cellcut/C02942B3.adjusted.gem \
--tif Alignment/ssDNA_C02942B3_regist_regist_transposed_aligned_by_Spateo.tiff \
--outdir Watershed --outprefix C02942B3

cellist watershed --gem /fs/home/dongxin/Projects/heart/202401_Stereo_seq/result_C02942B3/041.cellcut/C02942B3.adjusted.gem \
--tif Alignment/ssDNA_C02942B3_regist_regist_transposed_aligned_by_Spateo.tiff \
--outdir Watershed --outprefix C02942B3
'''

def CellposeParser(subparsers):
    parser = subparsers.add_parser("cellpose", 
        help = "Run initial cellpose segmentation on the staining image. ")
    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("--gem", dest = "spot_expr_file", default = None,
        help = "Bin1 gem file generated by SAW pipeline located in '041.cellcut' or '04.tissuecut'. ")
    group_input.add_argument("--tif", dest = "regist_tif_file", default = None,
        help = "Registered tif file generated by SAW pipeline located in '04.tissuecut' or the refined tif file generated by 'cellist align'. ")
    group_input.add_argument("--no-local-threshold", dest = "no_local_threshold", action = "store_true", 
        help = "Whether or not to perform local thresholding on the image. "
        "If not set, Cellist will perform local thresholding. If set, the local thresholding will be skipped. ")
    group_input.add_argument("--diameter", dest = "diameter", type = int, default = 15,
        help = "The average cell/nucleus diameter (in pixel) to run cellpose model. Default: 15. ")
    group_input.add_argument("--model-type", dest = "model_type", default = 'nuclei', choices = ["nuclei", "cyto", "cyto2", "cyto3"],
        help = "Cellpose model type. cyto: cytoplasm model; nuclei: nucleus model; "
        "cyto2: cytoplasm model with additional user images; "
        "cyto3: super-generalist model. Default: nuclei. ")
    group_input.add_argument("--flow-threshold", dest = "flow_threshold", type = float, default = 0.4,
        help = "The flow_threshold parameter in Cellpose is the maximum allowed error of the flows for each mask. "
        "Increase this threshold if cellpose is not returning as many ROIs as you’d expect. "
        "Similarly, decrease this threshold if cellpose is returning too many ill-shaped ROIs. Default: 0.4.")
    group_input.add_argument("--cellprob-threshold", dest = "cellprob_threshold", type = float, default = 0.0,
        help = "The cellprob_threshold parameter in Cellpose, varies from around -6 to +6. "
        "Decrease this threshold if cellpose is not returning as many ROIs as you’d expect. "
        "Similarly, increase this threshold if cellpose is returning too ROIs particularly from dim areas. Default: 0.0")
    group_input.add_argument("--expansion", dest = "expansion", action = "store_true", 
        help = "Whether or not to perform expansion after Cellpose segmentation. "
        "If set, expansion based on Cellpose segmentation will be performed. ")
    group_input.add_argument("--expansion-dist", dest = "expansion_dist", type = int, default = 8,
        help = "The expansion distance (in pixel). Default: 8. ")

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

def cellpose_seg(img, otsu_local_hole_rm, out_dir, out_img_prefix, diameter = 15, model_type = 'nuclei', flow_threshold = None, 
    cellprob_threshold = 0, expansion = False, expansion_dist = 8):
    logging.info("Cellpose segmentation... ")
    # 0. load model
    model = models.Cellpose(gpu=use_GPU, model_type = model_type, device = torch.device('cuda:1'))
    channels = [0, 0]
    segmented_cells, flow, style, diam = model.eval(otsu_local_hole_rm, diameter = diameter, flow_threshold = flow_threshold, channels=channels, cellprob_threshold = cellprob_threshold)
    logging.info("Number of segments: %s" %segmented_cells.max())
    cellpose_seg_file = os.path.join(out_dir, "%s_Cellpose_nucleus_matrix.npz" %out_img_prefix)
    segmented_cells_mat = csc_matrix(segmented_cells, dtype=np.int32)
    save_npz(cellpose_seg_file, segmented_cells_mat)
    fig, ax = plt.subplots(figsize=(20, 20))
    ax.imshow(mark_boundaries(img, segmented_cells, color = (1,0,0)))
    plt.savefig(os.path.join(out_dir, "%s_Cellpose_boundary.png" %out_img_prefix), bbox_inches='tight')
    props = regionprops_table(segmented_cells, 
                              properties=['label', 'area','centroid',
                                          'equivalent_diameter_area'])
    props_df = pd.DataFrame(props)
    props_df.to_csv(os.path.join(out_dir, "%s_Cellpose_nucleus_property.txt" %out_img_prefix), sep = "\t", index = False)
    if expansion:
        expanded = expand_labels(segmented_cells, distance = expansion_dist)
        cellpose_seg_file = os.path.join(out_dir, "%s_Cellpose_cells_expansion%s_matrix.npz" %(out_img_prefix, expansion_dist))
        expanded_mat = csc_matrix(expanded, dtype=np.int32)
        save_npz(cellpose_seg_file, expanded_mat)
        props_expanded = regionprops_table(expanded, 
                                  properties=['label', 'area','centroid',
                                              'equivalent_diameter_area'])
        props_df_expanded = pd.DataFrame(props_expanded)
        props_df_expanded.to_csv(os.path.join(out_dir, "%s_Cellpose_cells_expansion%s_property.txt" %(out_img_prefix, expansion_dist)), sep = "\t", index = False)
        return segmented_cells, expanded
    else:
        return segmented_cells

def Cellpose(platform, gem_path, img_path, out_dir, out_prefix, no_local_threshold, diameter, model_type, flow_threshold, cellprob_threshold, expansion, expansion_dist):
    if not out_prefix:
        file_name, file_extension = os.path.splitext(gem_path)
        base_name = os.path.basename(file_name)
        out_gem_prefix = base_name
        file_name, file_extension = os.path.splitext(img_path)
        base_name = os.path.basename(file_name)
        out_img_prefix = base_name
    else:
        out_gem_prefix = out_prefix
        out_img_prefix = out_prefix
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    # --------------------- read image file and thresholding --------------------- 
    threshold_file = os.path.join(out_dir, "%s_thresholding.tif" %out_img_prefix)
    if os.path.exists(threshold_file):
        otsu_local_hole_rm = imread(threshold_file)
        img = imread(img_path)
    else:
        img, otsu_local_hole_rm = Thresholding(platform, img_path, out_dir, out_img_prefix, no_local_threshold)
    # --------------------- watershed segmentation --------------------- 
    if expansion:
        segmented_cells, expanded = cellpose_seg(img, otsu_local_hole_rm, out_dir, out_img_prefix, diameter = diameter, model_type = model_type,
        flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold,
        expansion = expansion, expansion_dist = expansion_dist)
    else:
        segmented_cells = cellpose_seg(img, otsu_local_hole_rm, out_dir, out_img_prefix, diameter = diameter, model_type = model_type,
        flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold,
        expansion = expansion, expansion_dist = expansion_dist)
    # --------------------- read gem file --------------------- 
    logging.info("Reading count file...")
    gem_df = pd.read_csv(gem_path, sep = "\t", comment = "#")
    bin1_out_file = os.path.join(out_dir, "%s_bin1.h5" %out_gem_prefix)
    countname = gem_df.columns[3]
    gem_df = gem_to_mat(gem_df = gem_df, outfile = bin1_out_file, countname = countname)
    coord_df = gem_df.drop_duplicates(subset = ['x_y'])
    coord_df = coord_df[["x", "y", "x_y"]]
    # --------------------- write segmentation result --------------------- 
    if expansion:
        nucleus_barcode_df, expanded_barcode_df = write_segmentation_coord(segmented_cells, coord_df, out_dir, out_gem_prefix, seg_method = "Cellpose", expansion = expansion, expanded = expanded, expansion_dist = expansion_dist)
    else:
        nucleus_barcode_df = write_segmentation_coord(segmented_cells, coord_df, out_dir, out_gem_prefix, seg_method = "Cellpose")
        expanded_barcode_df = None
    write_segmentation_cell(nucleus_barcode_df, gem_df, out_gem_prefix, out_dir, seg_method = "Cellpose", expansion = expansion, expanded_barcode_df = expanded_barcode_df, expansion_dist = expansion_dist)
    logging.info("Drawing segmentation plot...")
    draw_segmentation(nucleus_barcode_df, seg_res = "Cellpose", out_prefix = "%s_Cellpose" %out_gem_prefix,
        out_dir = out_dir, x = "x", y = "y", figsize = (80, 80))
    if expansion:
        draw_segmentation(expanded_barcode_df, seg_res = 'Cellpose_expansion_%s' %(expansion_dist), out_prefix = "%s_Cellpose_expansion_%s" %(out_gem_prefix, expansion_dist),
            out_dir = out_dir, x = "x", y = "y", figsize = (80, 80))
    logging.info("Drawing segmentation plot is done.")
    logging.info("All done!")

if __name__ == '__main__':
    parser = WatershedParser()
    platform = parser.platform
    gem_path = parser.spot_expr_file
    img_path = parser.regist_tif_file
    out_dir = parser.out_dir
    out_prefix = parser.out_prefix
    no_local_threshold = parser.no_local_threshold
    diameter = parser.diameter
    model_type = parser.model_type
    flow_threshold = parser.flow_threshold
    cellprob_threshold = parser.cellprob_threshold
    expansion = parser.expansion
    expansion_dist = parser.expansion_dist
    Cellpose(platform = platform, gem_path = gem_path, img_path = img_path, out_dir = out_dir, out_prefix = out_prefix, no_local_threshold = no_local_threshold,
         diameter = diameter, model_type = model_type,flow_threshold = flow_threshold, cellprob_threshold = cellprob_threshold, expansion = expansion, expansion_dist = expansion_dist)
