# Author: Cem Emre Akbas
# e-mail: akbas@fi.muni.cz
# Center for Biomedical Image Analysis
# Faculty of Informatics, Masaryk University

from glob import glob
import os
import subprocess
from PIL import Image
import numpy as np
import time

start_time = time.time()
# This script finds and merges all available segmentation results of a chosen dataset.
# So, <dataset> parameter should be changed after getting the merging result.
# After getting the merging result, change the <dataset> and repeat the process above.

## -- Dataset name ------------------------- CHANGE AFTER EACH RUN -----------------------------------------
# Change the dataset name after each run
# Note that "Fluo-N3DH-CHO/01_RES" and "Fluo-N3DH-CHO/02_RES" are considered as different datasets in the concept of this script
dataset = "Fluo-N2DH-GOWT1/02_RES"

## -- Paths of Data and Executables ---- CHANGE ONLY BEFORE THE FIRST RUN ----------------------------------
# Executable paths: Download them from http://crl.med.harvard.edu/software/CRKIT/index.php. Change the paths below only before the first run
executable_path_crlConvertBetweenFileFormats = "C:/Computational_Radiology_Laboratory/CRKit/bin/crlConvertBetweenFileFormats.exe"     # should contain the full path of converter executable
executable_path_crlSTAPLE = "C:/Computational_Radiology_Laboratory/CRKit/bin/crlSTAPLE.exe"
executable_path_crlIndexOfMaxComponent = "C:/Computational_Radiology_Laboratory/CRKit/bin/crlIndexOfMaxComponent.exe"
executable_path_crlOverlapStats3d = "C:/Computational_Radiology_Laboratory/CRKit/bin/crlOverlapStats3d.exe"

# parent_directory should be the parent directory of all available results
# It should contain folders named as isbiXXX which respresent CTC participants
# Each isbiXXX folder contains submitted results by that participant in different subfolders that are named as <dataset>
# Change the parent_directory only before the first run
parent_directory = "C:/Users/ceakb/Documents/CTC_Training_Results/"

# goldGT_parentdir should contain the Ground Truth annotators for CTC Training data.
# Segmentation Gold GT datasets should be stored in <goldGT_parentdir>/<dataset>/01_GT/SEG or <goldGT_parentdir>/<dataset>/02_GT/SEG
goldGT_parentdir = "C:/Users/ceakb/Documents/CTC_TrainingGT_Annotators/"

## ---------------------------------------------------------------------------------------------------------
# Run the algorithm. Generated results are stored in <parent_directory> as seperate folders.

# This section finds all available segmentation results of a chosen dataset.
# Then it converts all the images into .nrrd image format which is required by STAPLE executables.
dirs = glob(parent_directory + "/*/")
print(dirs)

for d in dirs:      # Here we expect d to take values in the form of isbiXXX
    data_dir = d + "/" + dataset
    if (os.path.isdir(data_dir)):   # Check if that participant has results for this dataset
        os.chdir(data_dir)
        for x in os.listdir(data_dir):  # Go through all the .tif files in the folder
            if x.endswith(".tif"):
                im = Image.open(x)
                im.save(x, dpi=(25, 25))   # Fix dpi value for all available results. STAPLE executable requires inputs with same dpi
                subprocess.call(executable_path_crlConvertBetweenFileFormats + " -in " + os.getcwd() + "/" + x + " -out " + os.getcwd() + "/" + x.split(".")[0] + ".nrrd")

print("All images in the subdirectories of " + parent_directory + " containing " + dataset + " results are converted to nrrd format with fixed dpi values.")

# This section is the first step of the merging process. It creates a 4D probability map
# using all available segmentation results for each single frame.

output_dir = parent_directory + dataset.split("/")[0] + "_" + dataset.split("/")[1] + "_STAPLEoutputs/"
if (not os.path.isdir(output_dir)):
    os.makedirs(output_dir)     # create output directories if they don't exist

dirs = glob(parent_directory + "/*/")
print(dirs)
datadirs = []

for d in dirs:
    data_dir = d + "/" + dataset
    if (os.path.isdir(data_dir)):  # keep the folders that has results for the current dataset
        datadirs.append(data_dir)
        os.chdir(data_dir)
        filelist = [x for x in os.listdir(data_dir) if x.endswith(".nrrd")]     # keep the nrrd image file paths for processing below

for imname in filelist:
    imstr = ""
    for dd in datadirs:
        os.chdir(dd)
        impath = os.getcwd() + "/" + imname
        imstr = imstr + " " + impath    # imstr collects paths of all available segmentation results to be merged
        #print(imstr)
    subprocess.call(executable_path_crlSTAPLE + " -o " + output_dir + "weights_" + imname + " " + imstr)

print("4D Probability maps are created by combining all the results available for " + dataset)

# This section is the second step of the merging process. It converts 4d probability maps
# into segmentation masks which are our desired outputs.
input_dir = output_dir      # This folder should contain only 4D probability maps, weights_maskXXX.nrrd
output_dir_seg = parent_directory + dataset.split("/")[0] + "_" + dataset.split("/")[1] + "_STAPLEoutputs/seg/"

if (not os.path.isdir(output_dir_seg)):     # create output directories for segmentation masks
    os.makedirs(output_dir_seg)
os.chdir(input_dir)

for x in os.listdir(input_dir):
    if x.endswith(".nrrd"):
        subprocess.call(executable_path_crlIndexOfMaxComponent + " " + x + " " + output_dir_seg + "seg_" + x.split(".")[0].split("_")[1] + ".nrrd")    # call the merging process
        print("Merging is completed for image " + x)

# This section converts gold segmentation ground truth images into .nrrd format and saves them into the same folder
goldGT_dir = goldGT_parentdir + dataset.split("/")[0] + "/" + dataset.split("/")[1][0:3] +"GT/SEG/"
os.chdir(goldGT_dir)


for x in os.listdir(goldGT_dir):
    if x.endswith(".tif"):
        im = Image.open(x)
        im.save(x, dpi=(25, 25))   # Fix dpi values all images at the same value for the proper run of comparison executable
        subprocess.call(executable_path_crlConvertBetweenFileFormats + " -in " + x + " -out " + x.split(".")[0] + ".nrrd")

print("All images in the subdirectories of " + goldGT_dir + " results are converted to nrrd format with fixed dpi values.")

# This section compares STAPLE outputs with gold segmentation ground truth and produces the final result.
# Results are stored as a txt file in the folder of STAPLE outputs
results_directory = parent_directory + dataset.split("/")[0] + "_" + dataset.split("/")[1] + "_STAPLEoutputs/seg/"

output = ""
jaccard_array = np.array([])
for x in os.listdir(results_directory):
    if x.endswith(".nrrd"):
        gt_impath = goldGT_dir + "man_seg*" + x.split(".")[0][-3:] + "*.nrrd"
        result_impath = results_directory + x
        if glob(gt_impath).__len__() == 1:   # the exact match case
            if(os.path.isfile(str(glob(gt_impath)[0]))):     # for Real datasets, GT image is not available for all frames. Do comparison if GT exists.
                print("[INFO] Reading pair: " + str(glob(gt_impath)[0]))
                print("[INFO] Reading pair: " + result_impath)
                pipe = subprocess.Popen(executable_path_crlOverlapStats3d + " " + str(glob(gt_impath)[0]) + " " + result_impath + " 0", shell=True, stdout=subprocess.PIPE).stdout
                cmd_out = pipe.read()
                jaccard_pos = str.find(cmd_out.__str__(), 'jaccard')
                dice_pos = str.find(cmd_out.__str__(), 'dice')
                current_jaccard = float(cmd_out.__str__()[jaccard_pos+8:dice_pos-1])
                jaccard_array = np.append(jaccard_array, current_jaccard)
                output = output + "[INFO] Reading pair: " + str(glob(gt_impath)[0]) + " \n [INFO] Reading pair: " + result_impath + " \n" + cmd_out.__str__() + " \n \n"
                print()

elapsed_time = time.time() - start_time
output = output + "******** FINAL RESULT ******** \n The SEG score for the " + dataset + " dataset is: " + str(np.mean(jaccard_array)) + "\nComputation time is " + str(elapsed_time) + " seconds."
f = open(results_directory + 'results.txt', 'w')
f.write(output)
f.close()
print("Given segmentation results are compared with the gold GT. Results are saved in " + results_directory + "results_" + dataset.split("/")[0] + "-" + dataset.split("/")[1] + ".txt")