'''All samples were of cells in suspension, so their confluency on the slide
was less than 50%, resulting in considerable background noise, which
dramatically reduces the apparent mean cell signal. Without proper background
signal elimination, the quantification of the image would result in a gross
underestimation of the actual cell signal intensity. To resolve this issue, all
channel images were converted to an array of intensity in greyscale (pixel
values ranging between 0-255). 

Background determination: 1024*1024 pixel images were split into 400 square
blocks and the mean and standard deviation (SD) of the pixels in each block
were calculated. The means of blocks with SD<1 were compared to the means of
all the blocks, and those within 2 SD of the means of all the blocks were
defined as “background”. The frequency of intensity values of the “background”
blocks was counted, and all values accounting for more than 1% of the
background were dropped from the entire array. The mean value of the entire
array, after elimination of the background values, was calculated. Only images
with the same cutoff for background values were used, as different values cause
over or underestimation. 

Standardization across the slide: To account for random differences in image
intensity, all channels were standardized to DAPI, as the nuclear intensity is
not influenced by the treatment. All DAPI means (background eliminated) of all
fields of a single slide were averaged. For each channel image of any field on
a given slide, the mean channel value (after background elimination), was
standardized to the DAPI average of the slide. 

Nuclear localization: Following background elimination, the remaining
DAPI-positive pixels were averaged and the pixel locations were noted. After
background elimination, the β-catenin channel mean was calculated only for
values located at the DAPI-positive pixel locations. Finally, the
nuclear-localized β-catenin mean was divided by the DAPI mean.'''

import os
import re
import pandas as pd
import numpy as np
import cv2
from openpyxl import load_workbook

'''czi file names must be:      sample_protein_treatment_photo# 
    for example: sample5_bCAT_Wnt-3a_4
    
    all sections can contain any char, except "_", which counts as a sepataror
    photo# - is just a number after the "_"
    if the _protein_ name contains 'CAT' (not caps sensitive) the program with perform a nuclear localization for all channels to DAPI
    
    CZIs must be exported as:
        .tif OR .jpg
        CHECK - V - individual channel image
        CHECK - V - use channel names
    '''
'''the .py file sholud be in the folder of the folders of the images exported from ZEN.
the .py file should be run with IDLE directly from there, or in an IDE with the directory set as the folder of image folders.
the folder directory can't contain spaces!!! best to use in a desktop folder'
    '''

def exclude_num(img, num):
    ''' excludes all appearances of the num fro, the img array (-> nan)
    returns the "clean" img array and its mean'''
    img[img < num] = np.nan
    mean_intensity = np.nanmean(img)
    return img, mean_intensity

def count_n_in_array(array, n):
    '''counts how many times the number n appears in an array'''
    return array[array==n].shape[0]

df = pd.DataFrame(columns = ['sample', 'exp_protein', 'treatment', 'photo_repeat', 'fluorophore', 'mean_intensity'])
df_nuc_loc = pd.DataFrame(columns = ['sample', 'exp_protein', 'treatment', 'photo_repeat', 'fluorophore', 'mean_on_dapi', 'mean_on_dapi/DAPI'])


df_indiv_withouts = pd.DataFrame()

dirname = os.path.dirname(__file__)
for subdir, dirs, files in os.walk(dirname):
    if dirs==[]:
        
        list_last_background_above_1 = []

        b_cat_for_nuc_loc = []
        for file_name in files:
            print(file_name)
            if (file_name.endswith('.tif') or file_name.endswith('.jpg')) and file_name.count('_')==4:
                file_data = re.split('_', file_name)
                file_data[4] = re.sub('-(.)*', '', file_data[4])  #remove the -T1.tif from the end of the color name
                
                file_dir = os.path.join(subdir, file_name)
    
                img = cv2.imread(file_dir, 0)
                img = img.astype('float')
               
        ### CHECK BLOCKS
                ## split img array into 400 sub array squares = blocks
                list_blocks = []
                for b in np.array_split(img, 20):
                    list_blocks.append(np.array_split(b, 20, axis=1))
                list_blocks = [val for sublist in list_blocks for val in sublist]
                
                ## check std of each block. less then 1->likely_background. else->cell
                list_blocks_likely_background, list_blocks_cells = [], []
                for b in list_blocks:
                    if np.std(b) < 1:
                        list_blocks_likely_background.append(b)
                    else:
                        list_blocks_cells.append(b)
                
                ## calculate mean & std of all sub arrays
                list_array_means = []
                for b in list_blocks_likely_background:
                    one_mean = np.mean(b)
                    list_array_means.append(one_mean)
                array_mean = np.mean(list_array_means)
                array_std = np.std(list_array_means)
                         
                ## check if each subarray's mean is within 2 std of the total array avg
                ## if yes->background. else->likely_cell
                list_blocks_background, list_blocks_likely_cells = [], []         
                for b in list_blocks_likely_background:
                    if (array_mean-(2*array_std)) < np.mean(b) < (array_mean+(2*array_std)):
                        list_blocks_background.append(b)
                    else:
                        list_blocks_likely_cells.append(b)
                
                ## check how many times each num appears in the entire img
                df_count_appr = pd.DataFrame(index=range(0,256), columns=['count_entire', '%_entire',
                                                                          'count_background', '%_background',
                                                                          'count_likely_cells', '%_likely_cells',
                                                                          'count_cells', '%_cells'])
                                                                                      # 'count_likely_background', '%_likely_background',
    
                for i in range(0,256):
                    df_count_appr.loc[i, 'count_entire'] = count_n_in_array(img, i)
                
                df_count_appr['%_entire'] = df_count_appr['count_entire']/sum(df_count_appr['count_entire'])*100
                           
                #check for last intensity with a count, and remove everything that cames after
                max_pixel_intensity_img = df_count_appr[df_count_appr['count_entire'] != 0].index[-1]
                df_count_appr = df_count_appr.loc[0:max_pixel_intensity_img]
        
                #background
                df_b_count_background = pd.Series(0, index=df_count_appr.index)
                for b in list_blocks_background:
                    for i in df_b_count_background.index:
                        df_b_count_background[i] = df_b_count_background[i] + count_n_in_array(b, i)
                df_count_appr['count_background'] = df_b_count_background
                df_count_appr['%_background'] = df_count_appr['count_background']/sum(df_count_appr['count_background'])*100
                
                    #check tha last inensity for which the % of the background in above 1 (to reccomend excluding)
                try:
                    last_back_int = df_count_appr['%_background'][df_count_appr['%_background']>1].index[-1]
                    list_last_background_above_1.append([file_data[4], last_back_int])
                    ###########
                    
                    ### FOR EACH FLOUROPHORE OF EACH IMAGE THE BACKGROUND IS SET AND SUBTRACTED INDIVIDUALLY
                    img[img < last_back_int+1] = np.nan
                    mean_intensity = np.nanmean(img)
                    file_data.append(mean_intensity)
        
                    # append file_data as a row in df
                    df = df.append(pd.Series(file_data, index = df.columns), ignore_index=True)
                      
                    if re.search('cat', file_name, re.IGNORECASE):
                        b_cat_for_nuc_loc.append([file_data, pd.DataFrame(img)])

                except IndexError:
                    print('IndexError: no background identified')
                    pass
        
        for data, array in b_cat_for_nuc_loc:
            if data[4] == 'DAPI':
                print('nuc_loc:\t' + '_'.join(data[0:5]))
                blue_locs = np.argwhere(array.notnull().values).tolist()
                blue_mean = np.nanmean(array)
                to_append = [*data[:-1], blue_mean, np.nan]
                df_nuc_loc = df_nuc_loc.append(pd.Series(to_append, index = df_nuc_loc.columns), ignore_index=True)

        for data, array in b_cat_for_nuc_loc:
            if data[4] != 'DAPI':
                print('nuc_loc:\t' + '_'.join(data[0:5]))
                loc_values = []
                for loc in blue_locs:
                    loc_values.append(array.iloc[loc[0], loc[1]])
                color_mean = np.nanmean(loc_values)
                to_append = [*data[:-1], color_mean, np.nan]
                df_nuc_loc = df_nuc_loc.append(pd.Series(to_append, index = df_nuc_loc.columns), ignore_index=True)
        
        df_last_to_exclude = pd.DataFrame(list_last_background_above_1, columns = ['fluorophore', 'without'])
        df_last_to_exclude = df_last_to_exclude.set_index('fluorophore')
                
        # summary of last to exclude (one file for all folders in the dir)
        for flour in df_last_to_exclude.index:
            df_indiv_withouts.loc[flour, subdir.split('\\')[-1]] = df_last_to_exclude.loc[flour, 'without']
              
### Add a column of each intensity standartized to the mean dapi intensity of the slide.
df2 = df[(df['fluorophore']=='DAPI')].groupby(['sample' ,'exp_protein', 'treatment', 'fluorophore']).mean()
df2.reset_index(inplace=True)

for i in df.index:
    sample, exp_protein, treatment = df.loc[i, 'sample'], df.loc[i, 'exp_protein'], df.loc[i, 'treatment']
    dapi = df2['mean_intensity'][(df2['sample']==sample) & (df2['exp_protein']==exp_protein) & (df2['treatment']==treatment)].iloc[0]
    df.loc[i, 'to_mean_dapi'] = df.loc[i, 'mean_intensity']/dapi


sample_names = '_'.join(list(df['sample'].unique()))

#export summary df with stand. to mean slide dapi (withou the '0' labled images)
export_summary_dir = dirname + '\\mean_intensities_' + sample_names + '.xlsx'
writer = pd.ExcelWriter(export_summary_dir, engine='xlsxwriter')
df.to_excel(writer, sheet_name='Sheet1', index=False)
workbook  = writer.book
worksheet = writer.sheets['Sheet1']
worksheet.conditional_format('G2:G10000', {'type': '3_color_scale'})
writer.save()
#add filters to the header of mean_intensities.xlsx
wb = load_workbook(export_summary_dir)
ws = wb.active
ws.auto_filter.ref = ws.dimensions
wb.save(export_summary_dir)
wb.close()

#export summary df ind_values_to_exclude:
summary_ind_values_to_exclude_dir_name = '\\'.join(subdir.split('\\')[:-1]) + '\\CHECK_values_to_exclude_' + sample_names + '.xlsx'
writer = pd.ExcelWriter(summary_ind_values_to_exclude_dir_name, engine='xlsxwriter')
df_indiv_withouts.to_excel(writer, sheet_name='Sheet1')
workbook  = writer.book
worksheet = writer.sheets['Sheet1']
worksheet.conditional_format('B2:KO10', {'type': '3_color_scale'})
writer.save()

if not df_nuc_loc.empty:
    #df_nuc_loc divide all by mean dapi
    df_nuc_dapi_only = df_nuc_loc[df_nuc_loc['fluorophore']=='DAPI']
    
    for i in df_nuc_loc.index:
        sample, exp_protein, treatment, photo_repeat = df_nuc_loc.loc[i, 'sample'], df_nuc_loc.loc[i, 'exp_protein'], df_nuc_loc.loc[i, 'treatment'], df_nuc_loc.loc[i, 'photo_repeat']
        dapi = df_nuc_loc['mean_on_dapi'][(df_nuc_loc['sample']==sample) & (df_nuc_loc['exp_protein']==exp_protein) & (df_nuc_loc['treatment']==treatment) & (df_nuc_loc['photo_repeat']==photo_repeat)].iloc[0]
        df_nuc_loc.loc[i, 'mean_on_dapi/DAPI'] = df_nuc_loc.loc[i, 'mean_on_dapi']/dapi
    
    export_nuc_loc_dir = dirname + '\\nuc_loc_' + sample_names + '.xlsx'
    writer = pd.ExcelWriter(export_nuc_loc_dir, engine='xlsxwriter')
    df_nuc_loc.to_excel(writer, sheet_name='Sheet1', index=False)
    workbook  = writer.book
    worksheet = writer.sheets['Sheet1']
    worksheet.conditional_format('G2:G10000', {'type': '3_color_scale'})
    writer.save()
    #add filters to the header of mean_intensities.xlsx
    wb = load_workbook(export_nuc_loc_dir)
    ws = wb.active
    ws.auto_filter.ref = ws.dimensions
    wb.save(export_nuc_loc_dir)
    wb.close()
