# -*- coding: utf-8 -*-
# @Author: dongqing
# @Date:   2024-01-05 20:29:00
# @Last Modified by:   dongqing
# @Last Modified time: 2024-04-17 12:28:57


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

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 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 write_10X_h5, write_segmentation_h5, write_segmentation_cell_coord, gem_to_mat
from Cellist.Plot import *

'''
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 WatershedParser(subparsers):
    parser = subparsers.add_parser("watershed", 
        help = "Run initial watershed 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("--min-distance", dest = "min_distance", type = int, default = 6,
        help = "The minimal allowed distance (in pixel) separating peaks for watershed segmentation. "
        "Peaks are the local maxima in a region of 2 * min_distance + 1 (i.e. peaks are separated by at least min_distance). Default: 6. ")
    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_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 Thresholding(platform, img_path, out_dir, out_img_prefix, no_local_threshold = False):
    img = imread(img_path)
    logging.info("Read image done. ")
    # 0. multiotsu thresholding
    logging.info("Thresholding... ")
    threshold = threshold_multiotsu(img, 3)[0]
    mask = img > threshold
    img_mask = img.copy()
    img_mask[~mask] = 0
    if no_local_threshold: 
        otsu_local = img_mask.copy()
        otsu_local_hole_rm = remove_small_holes(otsu_local)
        fig, ax = plt.subplots(1, 2, figsize=(30, 17), sharex=True, sharey=True)
        ax[0].imshow(img, cmap="gray")
        ax[0].set_title("Original", fontsize = 20)
        ax[1].imshow(mask, cmap = "gray")
        ax[1].set_title('Mask', fontsize = 20)
        for a in ax.ravel():
            a.axis('off')
        plt.savefig(os.path.join(out_dir, "%s_thresholding.png" %out_img_prefix), bbox_inches='tight')
    else:
        # 1. local thresholding
        local_thresh = threshold_local(img, block_size = 51, offset=0)
        local = img > local_thresh
        # 2. combine two thresholding results
        mask_local = img_mask > local_thresh
        otsu_local = np.logical_and(local, mask_local)        
        # 3. remove small holes
        otsu_local_hole_rm = remove_small_holes(otsu_local)
        if platform == "imaging":
            otsu_local_hole_rm = mask.copy()
        fig, ax = plt.subplots(2, 2, figsize=(30, 34), sharex=True, sharey=True)
        ax[0, 0].imshow(img, cmap="gray")
        ax[0, 0].set_title("Original", fontsize = 20)
        ax[0, 1].imshow(mask, cmap = "gray")
        ax[0, 1].set_title('Mask', fontsize = 20)
        ax[1, 0].imshow(local, cmap = "gray")
        ax[1, 0].set_title('Local thresholding', fontsize = 20)
        ax[1, 1].imshow(otsu_local_hole_rm, cmap = "gray")
        ax[1, 1].set_title('Mask + local thresholding', fontsize = 20)
        for a in ax.ravel():
            a.axis('off')
        plt.savefig(os.path.join(out_dir, "%s_thresholding.png" %out_img_prefix), bbox_inches='tight')
    imsave(os.path.join(out_dir, "%s_thresholding.tif" %out_img_prefix), otsu_local_hole_rm)
    logging.info("Threshold done. ")
    return img, otsu_local_hole_rm

def watershed_seg(img, otsu_local_hole_rm, out_dir, out_img_prefix, min_distance = 6, expansion = False, expansion_dist = 8):
    logging.info("Wartershed segmentation... ")
    # 0. find local markers
    distance_local = ndi.distance_transform_edt(otsu_local_hole_rm)
    local_max_coords = peak_local_max(distance_local, min_distance = min_distance, exclude_border=False)
    local_max_mask = np.zeros(distance_local.shape, dtype=bool)
    local_max_mask[tuple(local_max_coords.T)] = True
    markers = label(local_max_mask)
    fig, ax = plt.subplots(figsize=(20, 20))
    ax.imshow(img, cmap='gray')
    ax.plot(local_max_coords[:,1], local_max_coords[:,0], 'r*', markersize = 2)
    ax.set_title('Markers (%s local maxima)' % local_max_coords.shape[0], fontsize = 20)
    plt.savefig(os.path.join(out_dir, "%s_markers.png" %out_img_prefix), bbox_inches='tight')
    # 1. watershed segmentation
    segmented_cells = watershed(-distance_local, markers, mask = otsu_local_hole_rm)
    logging.info("Watershed number of segments: %s" %(len(np.unique(segmented_cells))))
    watershed_seg_file = os.path.join(out_dir, "%s_watershed_nucleus_matrix.npz" %out_img_prefix)
    segmented_cells_mat = csc_matrix(segmented_cells, dtype=np.int32)
    save_npz(watershed_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_cell_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_watershed_nucleus_property.txt" %out_img_prefix), sep = "\t", index = False)
    if expansion:
        expanded = expand_labels(segmented_cells, distance = expansion_dist)
        watershed_seg_file = os.path.join(out_dir, "%s_watershed_cells_expansion%s_matrix.npz" %(out_img_prefix, expansion_dist))
        expanded_mat = csc_matrix(expanded, dtype=np.int32)
        save_npz(watershed_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_watershed_cells_expansion%s_property.txt" %(out_img_prefix, expansion_dist)), sep = "\t", index = False)
        return segmented_cells, expanded
    else:
        return segmented_cells

def write_segmentation_coord(segmented_cells, coord_df, out_dir, out_gem_prefix, expanded = None, seg_method = "nucleus", expansion_dist = 8):
    logging.info("Writing segmentation results... ")
    segmented_nucleus_df = pd.DataFrame(segmented_cells)
    segmented_nucleus_df["x"] = segmented_nucleus_df.index
    segmented_nucleus_long_df = pd.melt(segmented_nucleus_df, id_vars = "x", 
                                         value_vars = segmented_nucleus_df.columns[0:-1],
                                         var_name = "y",
                                         value_name = "Watershed")
    del segmented_nucleus_df
    gc.collect()
    segmented_nucleus_long_df["Nucleus"] = 0
    segmented_nucleus_long_df.loc[segmented_nucleus_long_df["Watershed"] > 0, "Nucleus"] = 1
    xmin,xmax = coord_df['x'].min(), coord_df['x'].max()
    ymin,ymax = coord_df['y'].min(), coord_df['y'].max()
    segmented_nucleus_long_df = segmented_nucleus_long_df.loc[(segmented_nucleus_long_df['x'] >= xmin) & 
                                                          (segmented_nucleus_long_df['x'] <= xmax) &
                                                          (segmented_nucleus_long_df['y'] >= ymin) &
                                                          (segmented_nucleus_long_df['y'] <= ymax),:]

    x_list = segmented_nucleus_long_df.loc[:,'x'].tolist()
    y_list = segmented_nucleus_long_df.loc[:,'y'].tolist()
    x_y = [f"{a}_{b}" for a, b in zip(x_list, y_list)]
    segmented_nucleus_long_df['x_y'] = x_y
    # segmented_nucleus_long_df['x_y'] = segmented_nucleus_long_df['x'].astype(str) + '_' + segmented_nucleus_long_df['y'].astype(str)
    coord_slice_list = slice_coord(coord_df, 5000, 0)
    segmented_nucleus_long_df_barcode_sub_list = []
    for coord_tuple in coord_slice_list:
        xmin_slice, xmax_slice, ymin_slice, ymax_slice = coord_tuple
        coord_df_sub = coord_df.loc[(coord_df['x'] >= xmin_slice) & 
                      (coord_df['x'] < xmax_slice) &
                      (coord_df['y'] >= ymin_slice) &
                      (coord_df['y'] < ymax_slice),:]
        segmented_nucleus_long_df_sub = segmented_nucleus_long_df.loc[(segmented_nucleus_long_df['x'] >= xmin_slice) & 
                      (segmented_nucleus_long_df['x'] < xmax_slice) &
                      (segmented_nucleus_long_df['y'] >= ymin_slice) &
                      (segmented_nucleus_long_df['y'] < ymax_slice),:]
        segmented_nucleus_long_df_barcode_sub = pd.merge(coord_df_sub, segmented_nucleus_long_df_sub[['x_y','Nucleus','Watershed']], how = 'left', on = 'x_y')
        segmented_nucleus_long_df_barcode_sub_list.append(segmented_nucleus_long_df_barcode_sub)
    segmented_nucleus_long_df_barcode = pd.concat(segmented_nucleus_long_df_barcode_sub_list)
    # segmented_nucleus_long_df_barcode = pd.merge(coord_df, segmented_nucleus_long_df[['x_y','Nucleus','Watershed']], how = 'left', on = 'x_y')
    segmented_nucleus_long_df_barcode = segmented_nucleus_long_df_barcode.sort_values(['x', 'y'])
    nucl_coord_out_file = os.path.join(out_dir, "%s_watershed_nucleus_coord.txt" %(out_gem_prefix))
    segmented_nucleus_long_df_barcode.to_csv(nucl_coord_out_file, sep = "\t", index = False)
    logging.info("Writing segmentation results is done.")
    if expanded:
        segmented_cells_df = pd.DataFrame(expanded)
        segmented_cells_df["x"] = segmented_cells_df.index
        segmented_cells_long_df = pd.melt(segmented_cells_df, id_vars = "x", 
                                         value_vars = segmented_cells_df.columns[0:-1],
                                         var_name = "y",
                                         value_name = "Expansion")
        segmented_cells_long_df["Nucleus"] = 0
        segmented_cells_long_df.loc[segmented_nucleus_long_df["Watershed"] > 0, "Nucleus"] = 1
        segmented_cells_long_df['x_y'] = segmented_cells_long_df['x'].astype(str) + '_' + segmented_cells_long_df['y'].astype(str)
        segmented_cells_long_df_barcode = pd.merge(coord_df, segmented_cells_long_df[['x_y','Nucleus','Expansion',]], how = 'left', on = 'x_y')
        segmented_cells_long_df_barcode = segmented_cells_long_df_barcode.sort_values(['x', 'y'])
        segmented_cells_long_df_barcode.loc[segmented_cells_long_df_barcode['Expansion'] == 0, 'Expansion'] = np.nan
        cell_coord_out_file = os.path.join(out_dir, "%s_watershed_cells_expansion%s_coord.txt" %(out_gem_prefix, expansion_dist))
        segmented_cells_long_df_barcode.to_csv(cell_coord_out_file, sep = "\t", index = False)
        return segmented_nucleus_long_df_barcode, segmented_cells_long_df_barcode
    else:
        return segmented_nucleus_long_df_barcode

def write_segmentation_cell(segmented_nucleus_long_df_barcode, count_df, out_gem_prefix, out_dir, segmented_cells_long_df_barcode = None):
    logging.info("Writing cell-level expression matrix... ")
    count_df_watershed = pd.merge(count_df, segmented_nucleus_long_df_barcode[['x_y', 'Watershed']], how = "right", on = 'x_y')
    count_name = count_df_watershed.columns[3]
    count_df_watershed = count_df_watershed.loc[count_df_watershed['Watershed'] != 0, :]
    write_segmentation_h5(count_df_seg = count_df_watershed, seg_res = "Watershed", out_prefix = "%s_waterhsed" %out_gem_prefix, 
        out_dir = out_dir, count_name = count_name)
    write_segmentation_cell_coord(coord_df_seg = segmented_nucleus_long_df_barcode, seg_res = 'Watershed', 
        out_prefix = "%s_waterhsed" %out_gem_prefix, out_dir = out_dir)
    if segmented_cells_long_df_barcode:
        count_df_expansion = pd.merge(count_df, segmented_cells_long_df_barcode[['x_y', 'Expansion']], how = "right", on = 'x_y')
        count_name = count_df_expansion.columns[3]
        write_segmentation_h5(count_df_seg = count_df_watershed, seg_res = "Expansion", out_prefix = "%s_expansion" %out_gem_prefix, 
            out_dir = out_dir, count_name = count_name)
        write_segmentation_cell_coord(coord_df_seg = segmented_cells_long_df_barcode, seg_res = 'Expansion', 
            out_prefix = "%s_expansion" %out_gem_prefix, out_dir = out_dir)
    logging.info("Writing cell-level expression matrix is done.")

def Watershed(platform, gem_path, img_path, out_dir, out_prefix, min_distance, no_local_threshold):
    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 --------------------- 
    segmented_cells = watershed_seg(img, otsu_local_hole_rm, out_dir, out_img_prefix, min_distance = min_distance, expansion = False, expansion_dist = 8)
    # --------------------- 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 --------------------- 
    segmented_nucleus_long_df_barcode = write_segmentation_coord(segmented_cells, coord_df, out_dir, out_gem_prefix)
    write_segmentation_cell(segmented_nucleus_long_df_barcode, gem_df, out_gem_prefix, out_dir)
    logging.info("Drawing segmentation plot...")
    draw_segmentation(segmented_nucleus_long_df_barcode, seg_res = "Watershed", out_prefix = "%s_watershed" %out_gem_prefix,
        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
    min_distance = parser.min_distance
    no_local_threshold = parser.no_local_threshold
    Watershed(platform = platform, gem_path = gem_path, img_path = img_path, out_dir = out_dir, out_prefix = out_prefix, min_distance = min_distance, no_local_threshold = no_local_threshold)
