###################################
# Katherina Hemmen ~ Core Unit Fluorescence Imaging ~ RVZ
# katherina.hemmen@uni-wuerzburg.de
###################################

# Export for pixel-wise analysis and heterogeneity analysis

# Input: ROI-masked intensity images (green prompt, red prompt, red delay)
# and mean tau green

# Output: Table of 256 x 256 = 65535 lines and the following columns:
# Y pixel - X pixel- Mean Tau(green) - Number of photons(green) - Number of photons(red prompt) 
# - Number of photons(red delay) - eff Stoichiometry(PIE) - SgSr - app FRET efficiency

# 1. Importing of all required functions
from skimage import io, img_as_ubyte
import glob
import os
import numpy as np

# 2. Data to be processed
path = 'C:/Users/Downloads/DATA/*.ptu'  
# All files within this folder in PTU format will be used to generate a list of filenames.

# Loop over all files in folder
for file in glob.glob(path):
    filename = os.path.abspath(file).split(".")[0]  # Get the filenames
    output_filename = filename + '.er4'  # How to save the results
    
    print('Processing....' + filename)

    # Donor mean fluorescence lifetime - based on symphotime export
    mean_tau_green_img = filename + '_tauFF.tif'
    mean_tau_green = io.imread(mean_tau_green_img)

    # Intensity images - split by color channel and time windows (prompt & delay)
    green_intensity_img = filename + '_gp.tif'
    green_intensity = io.imread(green_intensity_img)[0]
    
    red_prompt_intensity_img = filename + '_rp.tif'
    red_prompt_intensity = io.imread(red_prompt_intensity_img)[0]
    
    red_delay_intensity_img = filename + '_rd.tif'
    red_delay_intensity = io.imread(red_delay_intensity_img)[0]

    # 2. Loading the generated cell mask
    cell_ROI = filename + '_ROI.tif'
    mask = img_as_ubyte(io.imread(cell_ROI))

    # 3. Select the pixels of interest
    tau_green = mean_tau_green * mask
    nD = green_intensity * mask
    nA_FRET = red_prompt_intensity * mask
    nA_direct = red_delay_intensity * mask

    # 4. Calculate the derived parameter
    #----------------------------------------------
    # 4.1 effective Stoichiometry-PIE 
    # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # Note: effective Spie as we omit all correction factors

    Spie = (nA_FRET + nD) / (nA_FRET + nD + nA_direct)

    # 4.2 Green-Red intensity ratio, Sg/Sr 
    # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # ~> related also to FRET efficiency, but uncorrected intensities are used

    SgSr = nD / nA_FRET

    # 4.3 apparent FRET-efficiency (proximity ratio) based on intensities
    # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    # For correct FRET efficiency, a correction factor is required: gamma
    # Gamma defines the ratio of the detector sensitivities
    # E = nA(FRET) /(nA(FRET) + gamma * nD) <=> Eapp = nA(FRET) /(nA(FRET) + nD)

    Eapp = nA_FRET / (nA_FRET + nD)

    # 5. Reshape the parameter matrices
    # All parameters are reshaped from a 256 x 256 array -> 1 x 65'536 array
    lines, pixel = mean_tau_green.shape
    tau_green_1D = tau_green.reshape((lines * pixel))
    nD_1D = nD.reshape((lines * pixel))
    nA_FRET_1D = nA_FRET.reshape((lines * pixel))
    nA_direct_1D = nA_direct.reshape((lines * pixel))
    Spie_1D = Spie.reshape((lines * pixel))
    SgSr_1D = SgSr.reshape((lines * pixel))
    Eapp_1D = Eapp.reshape((lines * pixel))

    # 6. Generate a list of pixel coordinates
    pixel_id = np.indices((lines, pixel))
    pixel_id_1D = pixel_id.reshape((2, lines*pixel))
    Y_pixel = np.array(pixel_id_1D[0, :])
    X_pixel = pixel_id_1D[1, :]
    
    # 7. Export results as text files
    # These files can be read by any text editor or e.g. Margarita from the Seidel-Software Package
    header = "Y pixel\tX pixel\tMean Tau(green)\tNumber of photons(green)\tNumber of photons(red prompt)\tNumber of photons(red delay)\teff Stoichiometry(PIE)\tSgSr\tapp FRET efficiency"
    
    # Change the saving directory!
    np.savetxt(
        output_filename,
        np.vstack([Y_pixel, X_pixel, tau_green_1D, nD_1D, nA_FRET_1D, nA_direct_1D, Spie_1D, SgSr_1D, Eapp_1D]).T,
        delimiter='\t',
        header=header
        )