"""
Copyright by Jan-Philipp_Praetorius

Research Group Applied Systems Biology - Head: Prof. Dr. Marc Thilo
Figge
https://www.leibniz-hki.de/en/applied-systems-biology.html
HKI-Center for Systems Biology of Infection
Leibniz Institute for Natural Product Research and Infection Biology -
Hans Knöll Insitute (HKI)
Adolf-Reichwein-Straße 23, 07745 Jena, Germany
"""

"""
    
    resolution = 0.2533
    # A_min_um = 15um (lower-bound) <=> A_max_um = 34um (upper-bound) for yeast cells
    A_min = 60 # 15um / 0.2533 um = 60 px      
    A_max = 134 # 34um / 0.2533 um = 134 px
   
    # convert areas-sizes from pixel to microns (discard objects bigger than 2000um for the plotting purpose)
    areas_arr_um = areas_arr * resolution
    yeast_areas_um = yeast_areas * resolution
    hype_areas_um = hype_areas * resolution
    hype_areas_um = hype_areas_um[hype_areas_um < 2000]

"""

import os
import numpy as np
from glob import glob
import matplotlib.pyplot as plt
import cv2
from skimage import io, filters, util, morphology, measure
from PIL import Image
from scipy import ndimage
import pandas as pd
from scipy import stats


def plot_img(img, mycmap=None, savefig=None):

    plt.figure(figsize=(24,20))
    if mycmap is None:
        plt.imshow(img)
    else:
        plt.imshow(img, cmap=mycmap)
    plt.grid(True)
    x_ticks = np.arange(0, img.shape[1], step=3000)
    plt.xticks(x_ticks, fontsize=10)
    y_ticks = np.arange(0, img.shape[0], step=3000)
    plt.yticks(y_ticks, fontsize=10)
    
    if savefig is not None:
        plt.savefig(savefig)
        
    plt.show()
    print('img-shape and type:', img.shape, img.dtype, 'min-max', img.min(), img.max())
    
    
def segment_kidney(img):
    
    print(f'[segment kidney] shape of input image: {img.shape}')
        
    # (1) convert to grayscale image
    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)    
    print(f'[segment kidney] convert to grayscale image with shape: {img_gray.shape}')
    
    # (2) binarize image via otsu thresholding
    threshold = filters.threshold_otsu(img_gray)
    img_binary = img_gray > threshold
    print(f'[segment kidney] perform otsu-thresolding with threshold: {threshold}')
    
    # (3) invert the intensity
    img_binary = util.img_as_ubyte(util.invert(img_binary))
    print(f'[segment kidney] invert the value - resulting type: {img_binary.dtype} and shape: {img_binary.shape}')
    
    # (4) perform morphology closing with disk and radius=7 
    kernel = morphology.disk(7)
    img_binary = cv2.morphologyEx(img_binary, cv2.MORPH_CLOSE, kernel)
    print(f'[segment kidney] performed morphology closing with disk-shape: {kernel.shape}')
        
    # (5) remove outliers noise via replaces a pixel by the median of the pixels in the surrounding if it deviates from the median by more than a certain value (the threshold)
    # https://imagej.nih.gov/ij/docs/menus/process.html#noise
    # img_binary = filters.rank.noise_filter(img_binary, morphology.disk(13)) # via skimage 
    img_binary = cv2.medianBlur(img_binary, 13)
    print(f'[segment kidney] remove outliers noise with disk-shape: {kernel.shape}')
        
    # (6) invert the intensity for subsequent area opening
    img_binary = util.invert(img_binary)
    print(f'[segment kidney] invert the value - resulting type: {img_binary.dtype} and shape: {img_binary.shape}')
    
    # (7) apply an area opening with a area threshold of 150000
    min_size = 15000
    
    ### opencv    
    #nb_blobs, im_with_separated_blobs, stats, _ = cv2.connectedComponentsWithStats(img_binary)
    ## here, we're interested only in the size of the blobs, contained in the last column of stats.
    #sizes = stats[:, -1]
    ## the following lines result in taking out the background which is also considered a component, which I find for most applications to not be the expected output.
    ## you may also keep the results as they are by commenting out the following lines. You'll have to update the ranges in the for loop below. 
    #sizes = sizes[1:]
    #nb_blobs -= 1

    ## minimum size of particles we want to keep (number of pixels). Here, it's a fixed value, but you can set it as you want, eg the mean of the sizes or whatever.
    ## output image with only the kept components
    #img_result = np.zeros(shape=(img_binary.shape), dtype=np.uint8)
    ## for every component in the image, keep it only if it's above min_size
    #for blob in range(nb_blobs):
    #    if sizes[blob] >= min_size:
    #        # see description of im_with_separated_blobs above
    #        img_result[im_with_separated_blobs == blob + 1] = 255
    #print(f'[segment kidney] apply an area opening with a area threshold of {min_size} - resulting type: {img_result.dtype} and shape: {img_result.shape}')
    
    ### skimage
    img_binary = morphology.area_opening(img_binary, area_threshold=min_size, connectivity=1)
    #img_binary = morphology.remove_small_objects(img_binary, min_size=min_size, connectivity=1)
    print(f'[segment kidney] apply an area opening with a area threshold of {min_size} - resulting type: {img_binary.dtype} and shape: {img_binary.shape}')
    
    # (8) invert the intensity for final binary mask
    img_binary = util.invert(img_binary)
    #img_binary = util.invert(img_result)
    print(f'[segment kidney] invert the value')
    
    # plot result
    plot_img(img_binary, mycmap='gray')
    
    return img_binary


def correct_kindey_mask(path_img, idx_ROI_remove, save_path):

    Image.MAX_IMAGE_PIXELS = 5000000000
    
    img_mask = io.imread(path_img)
    print(f'[correct_kindey_mask] image shape: {img_mask.shape} - type: {img_mask.dtype}')
    
    img_label = measure.label(img_mask, connectivity=img_mask.ndim)
    print(f'[correct_kindey_mask] perform label image')
        
    props = measure.regionprops_table(img_label, properties=['label', 'bbox', 'area', 'centroid'])
    print(f'[correct_kindey_mask] measure all ROIs properties')
        
    df = pd.DataFrame(props)
    df = df.sort_values(by=['area'], ascending=False)

    label_to_be_removed_ROIs = list(df.iloc[idx_ROI_remove]['label'])
    print(f'[correct_kindey_mask] labels to be removed: {label_to_be_removed_ROIs}')

    for label_tmp in label_to_be_removed_ROIs:
        img_label[img_label==label_tmp] = 0
        print(f'[correct_kindey_mask] set label: {label_tmp} to 0')

    img_result = img_label > 0
    img_result = util.img_as_ubyte(img_result)

    img_filter_out = img_mask - img_result

    # save image 
    path_save_corrected_kidney_mask = os.path.join(save_path, 'kidney_mask', os.path.basename(path_img).split('.')[0]+'_corrected.png')
    io.imsave(path_save_corrected_kidney_mask, img_result)
    print(f'[correct_kindey_mask] saved image to: {path_save_corrected_kidney_mask}')

    # plot the result
    fig, ax1 = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(30, 30))
    ax1.imshow(img_result, cmap='gray')
    ax1.contour(img_filter_out.astype(np.uint8), colors='red', linewidths=3)
    ax1.set_title("Origin image", fontsize=30)
    plt.savefig(os.path.join(save_path, 'kidney_mask', 'CONTROL', os.path.basename(path_img).split('.')[0]+'_corrected.png'))
    plt.show()
    
    return df


def measure_ROIs(path_directory, save_path, path_directory_overlay=None):
    
    Image.MAX_IMAGE_PIXELS = 5000000000
    
    filenames = glob(path_directory)
    print(f'[measure_ROIs] detected files:\n', *[x.split('/')[-1] for x in np.sort(filenames)], sep='\n')
    
    if path_directory_overlay is not None:
        filenames_overlay = glob(path_directory_overlay)
        print(f'[measure_ROIs] detected files:\n', *[x.split('/')[-1] for x in np.sort(filenames_overlay)], sep='\n')    
    
    df_collected = pd.DataFrame()
    
    for idx, filename in enumerate(filenames):
        
        sample = (os.path.basename(filename).split('.')[0]).replace("_corrected", "")
        
        # combine filenames if necessary:
        if path_directory_overlay is not None:
            filename_overlay = [x for x in filenames_overlay if sample in x][0]
            sample_overlay = (os.path.basename(filename_overlay).split('.')[0]).replace("_corrected", "")
                
        # read image(s)
        img = io.imread(filename)
        print(f'\n[measure_ROIs] [{idx}/{len(filenames)-1}] read file: {filename} with image shape: {img.shape} and sample-name: {sample}')
        
        if path_directory_overlay is not None:
            img_overlay = io.imread(filename_overlay)            
            print(f'[measure_ROIs] [{idx}/{len(filenames_overlay)-1}] read overlay-file: {filename_overlay} with image shape: {img_overlay.shape} and sample-name: {sample_overlay}')
        
            # bitwise overlay (AND)
            img = cv2.bitwise_and(img, img_overlay)      
            print(f'[measure_ROIs] overlay performed')
            
        # perform the ROI labeling
        img_label = measure.label(img, connectivity=img.ndim) # 56 sec on <V_210_20_B_Grocott_cutOut>
        #img_label = ndimage.label(img, structure=np.ones((3, 3)))[0] # 31 sec on <V_210_20_B_Grocott_cutOut>
        print(f'[measure_ROIs] label image performed with skimage.measure.label(img, connectivity=img.ndim): {img.ndim}')
        
        ### skimage-version: 0.19.X
        #myproperties = ['num_pixels', 'area', 'area_bbox', 'area_convex', 'area_filled', 'axis_major_length', 'axis_minor_length', 'euler_number', 'feret_diameter_max', 'label', 'perimeter', 'perimeter_crofton', 'solidity']
        ### skimage-version: 0.18.X
#         myproperties = ['area', 'bbox_area', 'convex_area', 'euler_number', 'feret_diameter_max', 'filled_area', 'label', 'major_axis_length', 'minor_axis_length', 'perimeter', 'perimeter_crofton', 'solidity']
        myproperties = ['area', 'convex_area', 'feret_diameter_max', 'filled_area', 'label', 'perimeter', 'perimeter_crofton'] # filter out some to speed-up the quantification of detected ROIs
              
        props_sample = measure.regionprops_table(img_label, properties=myproperties)
        
        df_sample = (
            pd.DataFrame(props_sample)
            .rename(columns={"area": "Area"})
            .assign(Sample=sample)
            .assign(Circularity=lambda df_: (4 * np.pi * df_.Area) / (df_.perimeter_crofton * df_.perimeter_crofton))
        )
        print(f'[measure_ROIs] measure all ROIs properties with shape: {df_sample.shape} and table columns: {list(df_sample.columns)}')
                
        save_path_df_sample = os.path.join(save_path, os.path.basename(filename).split('.')[0] + '.csv')
        df_sample.to_csv(save_path_df_sample)
        print(f'[measure_ROIs] save DataFrame to: {save_path_df_sample}')
        
        df_collected = pd.concat([df_collected, df_sample])
        
#         return df_collected # TMP
    
    # save all ROIs per sample
    df_collected.to_csv(os.path.join(save_path, "all_samples.csv"))
        
    return df_collected


def apply_operator(path_filenames, save_figures, save_binary, MinArea, MaxSolidity):
        
    filenames = np.sort(glob(path_filenames))
    
    df_collected = pd.DataFrame()
    
    # iterate over each sample
    for idx, filename in enumerate(filenames):
        
        sample = (os.path.basename(filename).split('.')[0]).replace("_corrected", "")
        
        print(f'\n[apply_operator] [{idx}/{len(filenames)-1}] read sample: {sample} --- with corresponding original: {os.path.basename(filename)}')
        
        img_binary = io.imread(filename)
        
        # perform the ROI labeling
        img_label = measure.label(img_binary, connectivity=img_binary.ndim) # 56 sec on <V_210_20_B_Grocott_cutOut>
        #img_label = ndimage.label(img, structure=np.ones((3, 3)))[0] # 31 sec on <V_210_20_B_Grocott_cutOut>
        print(f'[apply_operator] label image performed with skimage.measure.label(img, connectivity=img.ndim): {img_binary.ndim}')
        
        ### skimage-version: 0.19.X
        #myproperties = ['num_pixels', 'area', 'area_bbox', 'area_convex', 'area_filled', 'axis_major_length', 'axis_minor_length', 'euler_number', 'feret_diameter_max', 'label', 'perimeter', 'perimeter_crofton', 'solidity']
        ### skimage-version: 0.18.X
#         myproperties = ['area', 'bbox_area', 'convex_area', 'euler_number', 'feret_diameter_max', 'filled_area', 'label', 'major_axis_length', 'minor_axis_length', 'perimeter', 'perimeter_crofton', 'solidity']
        myproperties = ['area', 'convex_area', 'feret_diameter_max', 'filled_area', 'label', 'perimeter', 'perimeter_crofton'] # filter out some to speed-up the quantification of detected ROIs
        
        props_sample = measure.regionprops_table(img_label, properties=myproperties)
        
        df_sample = (
            pd.DataFrame(props_sample)
            .rename(columns={"area": "Area"})
            .assign(Sample=sample)
            .assign(Circularity=lambda df_: (4 * np.pi * df_.Area) / (df_.perimeter_crofton * df_.perimeter_crofton))
            .assign(Solidity=lambda df_: df_.Area / df_.convex_area)            
        )
        print(f'[apply_operator] measure all ROIs properties DataFrame:\n {df_sample.sort_values(by=["Area"], ascending=False).head(20)}')
        
        # iterate over properties to filter outliers
        for index, row in df_sample.iterrows():
            
            filter_label_out = row['label']
            
            if row['Area'] > MinArea and row['Solidity'] < MaxSolidity:
                
                img_label[img_label == filter_label_out] = 0    
                print(f"[apply_operator] Filter out label: {filter_label_out}")
                
        img_binary = img_label > 0
        
        img_binary = util.img_as_ubyte(img_binary)
        
        save_plot = os.path.join(save_figures, sample+'.jpg')
        plot_img(img_binary, mycmap='gray', savefig=save_plot)
        
        save_filename = os.path.join(save_binary, sample+'.png')
        io.imsave(save_filename, img_binary)

        print(f'[FINISH] save binary file to: {save_filename}')
        
        df_collected = pd.concat([df_collected, df_sample])
        
    return df_collected


def add_p_value_annotation(fig, array_columns, subplot=None, _format=dict(interline=0.07, text_height=1.07, color='black'), test_to_perform='t-test'):
    ''' Adds notations giving the p-value between two box plot data (t-test two-sided comparison)
    
    Parameters:
    ----------
    fig: figure
        plotly boxplot figure
    array_columns: np.array
        array of which columns to compare 
        e.g.: [[0,1], [1,2]] compares column 0 with 1 and 1 with 2
    subplot: None or int
        specifies if the figures has subplots and what subplot to add the notation to
    _format: dict
        format characteristics for the lines
        
    https://stackoverflow.com/questions/67505252/plotly-box-p-value-significant-annotation
    
    use via: fig = add_p_value_annotation(fig, [[0,1], [0,2], [0,3]])

    Returns:
    -------
    fig: figure
        figure with the added notation
    '''
    # Specify in what y_range to plot for each pair of columns
    y_range = np.zeros([len(array_columns), 2])
    for i in range(len(array_columns)):
        y_range[i] = [1.01+i*_format['interline'], 1.02+i*_format['interline']]

    # Get values from figure
    fig_dict = fig.to_dict()

    # Get indices if working with subplots
    if subplot:
        if subplot == 1:
            subplot_str = ''
        else:
            subplot_str =str(subplot)
        indices = [] #Change the box index to the indices of the data for that subplot
        for index, data in enumerate(fig_dict['data']):
            #print(index, data['xaxis'], 'x' + subplot_str)
            if data['xaxis'] == 'x' + subplot_str:
                indices = np.append(indices, index)
        indices = [int(i) for i in indices]
        print((indices))
    else:
        subplot_str = ''

    # Print the p-values
    for index, column_pair in enumerate(array_columns):
        if subplot:
            data_pair = [indices[column_pair[0]], indices[column_pair[1]]]
        else:
            data_pair = column_pair

        # Mare sure it is selecting the data and subplot you want
        print('0:', fig_dict['data'][data_pair[0]]['name'], fig_dict['data'][data_pair[0]]['xaxis'])
        print('1:', fig_dict['data'][data_pair[1]]['name'], fig_dict['data'][data_pair[1]]['xaxis'])

        # Get the p-value
        if test_to_perform == "t-test":
            pvalue = stats.ttest_ind(
                fig_dict['data'][data_pair[0]]['y'],
                fig_dict['data'][data_pair[1]]['y'],
                equal_var=False,
            )[1]
        else:
            pvalue = stats.wilcoxon(
                fig_dict['data'][data_pair[0]]['y'],
                fig_dict['data'][data_pair[1]]['y'],
                equal_var=False,
            )[1]            
            
        if pvalue >= 0.05:
            symbol = 'ns'
        elif pvalue >= 0.01: 
            symbol = '*'
        elif pvalue >= 0.001:
            symbol = '**'
        else:
            symbol = '***'
        # Vertical line
        fig.add_shape(type="line",
            xref="x"+subplot_str, yref="y"+subplot_str+" domain",
            x0=column_pair[0], y0=y_range[index][0], 
            x1=column_pair[0], y1=y_range[index][1],
            line=dict(color=_format['color'], width=2,)
        )
        # Horizontal line
        fig.add_shape(type="line",
            xref="x"+subplot_str, yref="y"+subplot_str+" domain",
            x0=column_pair[0], y0=y_range[index][1], 
            x1=column_pair[1], y1=y_range[index][1],
            line=dict(color=_format['color'], width=2,)
        )
        # Vertical line
        fig.add_shape(type="line",
            xref="x"+subplot_str, yref="y"+subplot_str+" domain",
            x0=column_pair[1], y0=y_range[index][0], 
            x1=column_pair[1], y1=y_range[index][1],
            line=dict(color=_format['color'], width=2,)
        )
        ## add text at the correct x, y coordinates
        ## for bars, there is a direct mapping from the bar number to 0, 1, 2...
        fig.add_annotation(dict(font=dict(color=_format['color'],size=14),
            x=(column_pair[0] + column_pair[1])/2,
            y=y_range[index][1]*_format['text_height'],
            showarrow=False,
            text=symbol,
            textangle=0,
            xref="x"+subplot_str,
            yref="y"+subplot_str+" domain"
        ))
    return fig
        

def evaluate_ilastik(path, sigma=None):

    filenames = glob(path)

    df_all = pd.DataFrame()

    for idx, filename in enumerate(filenames):

        sample = os.path.basename(filename).split('=')[1].split('.')[0]
        
        df_result = pd.read_csv(filename)
        df_result['sample'] = sample
        if sigma is not None:
            df_result['sigma'] = sigma
        
        print(f'[evaluate_ilastik] [{idx}/{len(filenames)-1}] read sample: {sample} with shape: {df_result.shape}')
        
        df_all = pd.concat([df_all, df_result])

    return df_all


def binarize_probability_maps(path_filenames, save_figures, save_binary, thresholding_method):
    
    if thresholding_method not in ["otsu", "isodata"]:
        print(f'[binarize_probability_maps] <thresholding_method> is not a valid thresholding method, valid ones are : {["otsu", "isodata"]}') 
        return
    
    filenames = np.sort(glob(path_filenames))
    
    for idx, filename in enumerate(filenames):
        
        print(f'[binarize_probability_maps] [{idx}/{len(filenames)-1}] read sample: {filename} with corresponding original: {os.path.basename(filename)}')
        
        img_prob = io.imread(filename)
        
        if thresholding_method == "otsu":
            thresh = filters.threshold_otsu(img_prob)        
        elif thresholding_method == "isodata":
            thresh = filters.threshold_isodata(img_prob)        
        else:
            return
            
        img_binary = img_prob > thresh
        
        img_binary = util.img_as_ubyte(img_binary)
        
        save_plot = os.path.join(save_figures, os.path.basename(filename).split('tif')[0]+'jpg')
        plot_img(img_binary, mycmap='gray', savefig=save_plot)
        
        save_filename = os.path.join(save_binary, os.path.basename(filename).split('tif')[0]+'png')
        io.imsave(save_filename, img_binary)

        print(f'[FINISH] threshold: {thresh} with filename: {save_filename}')

        
def read_ometiff(input_path):
    
    # from apeer_ometiff_library import io, processing, omexmlClass

    with tifffile.TiffFile(input_path) as tif:
        array = tif.asarray()
        omexml_string = tif[0].image_description.decode('utf-8')

    # Turn Ome XML String to an Bioformats object for parsing
    metadata = omexmlClass.OMEXML(omexml_string)

    # Parse pixel sizes
    pixels = metadata.image(0).Pixels
    size_c = pixels.SizeC
    size_t = pixels.SizeT
    size_z = pixels.SizeZ
    size_x = pixels.SizeX
    size_y = pixels.SizeY

    # Expand image array to 5D of order (T, Z, C, X, Y)
    if size_c == 1:
        array = np.expand_dims(array, axis=-3)
    if size_z == 1:
        array = np.expand_dims(array, axis=-4)
    if size_t == 1:
        array = np.expand_dims(array, axis=-5)

    return array, omexml_string