#@ File (label = "Input directory (leave blank for processing front-most image)", style = "directory",  required=false) src_dir
#@ String (label = "Save location", choices={"Parent folder", "Same folder"}) save_location
#@ Boolean (label = "Z-Project") z_project
#@ String (label = "Channels to display separatedly in grays") separated_channels
#@ Integer (label = "Length of scale bar in µm (0 to disable)", min = 0) scale_length
#@ Float (label = "Pixel size im µm (0 to use existing calibration)", min = 0) pixel_size
#@ Boolean (label = "Process as film") process_as_film

import os
import re

from ij import IJ, ImagePlus, ImageStack
from ij.gui import NonBlockingGenericDialog
from ij.plugin import HyperStackConverter, MontageMaker, ChannelArranger, ZProjector
from ij.plugin.frame import RoiManager

# Fonction appelée pour chaque image trouvée dans le dossier
def processImage(filepath=None, output_basename=None):
    # Si l'image à traiter est déjà ouverte, on récupère son dossier
    if process_active_image:
        imp_original = IJ.getImage()
        image_dir = IJ.getDir("image")
        imp_original.hide()

        if not image_dir:
            image_dir = src_dir.getAbsolutePath()
        output_dir = os.path.dirname(image_dir) if save_location == "Parent folder" else image_dir
        output_basename = os.path.join(output_dir, os.path.splitext(imp_original.getTitle())[0])

        imp = imp_original.duplicate()

    # Sinon, on ouvre l'image dans le dossier indiqué
    else:
        imp = IJ.openImage(filepath)
        imp.hide()

        # On applique éventuellement les valeurs d'affichage précédentes, ou à défaut les valeurs extrémales
        for i in range(imp.getDimensions()[2]):
            imp.setC(i + 1)
            if i in mins_and_maxs:
                imp.setDisplayRange(mins_and_maxs[i][0], mins_and_maxs[i][1])
            else:
                IJ.resetMinAndMax(imp)
        imp.setC(1)

        # Si l'image est un z-stack, on se place au milieu
        if imp.getDimensions()[3] > 1:
            imp.setZ((imp.getDimensions()[3] + 1) / 2)

    if z_project and imp.getDimensions()[3] > 1:
        imp2 = ZProjector.run(imp, "max all")
        imp.close()
        imp = imp2

    # if (not imp.isHyperStack()) and imp.getStackSize() > 1:
    #     imp2 = HyperStackConverter.toHyperStack(imp, imp.getStackSize(), 1, 1)
    #     imp.changes = False
    #     imp.close()
    #     imp = imp2

    imp.show()
    IJ.selectWindow(imp.getID())

    # On demande à l'utilisateur de sélectionner des ROIs rectangulaires autour des cellules d'intérêt
    IJ.setTool("rectangle")
    rm = RoiManager.getRoiManager()
    rm.runCommand(imp, "Show All with labels")

    global roi_popup_location

    gui = NonBlockingGenericDialog("ROI selection")
    gui.addMessage(
        "Select region of interest (or add many regions with \"t\") \n Warning : Dont' forget to click before selecting, when using Shift-selection")
    gui.addCheckbox("Skip", False)
    if roi_popup_location is not None:
        gui.setLocation(roi_popup_location.x, roi_popup_location.y)
    gui.showDialog()

    roi_popup_location = gui.getLocation()

    # Si on ferme la fenêtre de dialogue, le retour de la fonction est "True" pour gérer l'annulation du script
    if not gui.wasOKed():
        return True

    # Si l'option "Skip" a été cochée, on passe à la prochaine image sans traiter l'image en cours
    if gui.getNextBoolean():
        imp.close()
        return False

    rois = rm.getRoisAsArray()
    rm.reset()
    rm.close()

    # Si aucun ROI n'a été ajouté au RoiManager ...
    if not rois:
        if imp.getRoi() is not None:
            IJ.run(imp, "Crop", "")  # Si il y a une sélection, on rogne, sinon on utilise l'image entière
        processImageCropped(imp, output_basename)

    # Sinon, on crée un tableau d'images en utilisant le RoiManager
    else:
        cropped_array = imp.crop(rois, "stack")
        imp.close()

        for i, cropped in enumerate(cropped_array):
            cropped.show()
            if cropped.getDimensions()[3] > 1:
                cropped.setZ((cropped.getDimensions()[3] + 1) / 2)
            interrupt = processImageCropped(cropped, output_basename + "_" + str(i + 1))
            if interrupt:
                return True

    if process_active_image:
        imp_original.show()

    return False


# Fonction appelée pour chaque ROI sélectionnée sur l'image source
def processImageCropped(imp, output_basename):

    # On demande à régler les niveaux de couleur à l'utilisateur
    imp.setDisplayMode(IJ.COLOR)
    IJ.selectWindow(imp.getID())
    IJ.run("Brightness/Contrast...")

    global threshold_popup_location

    gui = NonBlockingGenericDialog("Thresholding")
    gui.addMessage("Set your thresholds, then click OK")
    if threshold_popup_location is not None:
        gui.setLocation(threshold_popup_location.x, threshold_popup_location.y)
    gui.showDialog()

    threshold_popup_location = gui.getLocation()

    # Si on ferme la fenêtre de dialogue, le retour de la fonction est "True" pour gérer l'annulation du script
    if not gui.wasOKed():
        return True

    imp.hide()

    # On sauvegarde les nouvelles valeurs d'affichage :
    for i in range(imp.getDimensions()[2]):
        imp.setC(i + 1)
        mins_and_maxs[i] = (imp.getDisplayRangeMin(), imp.getDisplayRangeMax())
    imp.setC(1)

    # Si on a encore un stack (pas de Z-Project ni de film) on garde la slice actuelle uniquement
    if not process_as_film and imp.getDimensions()[3] > 1:
        IJ.run(imp, "Reduce Dimensionality...", "channels bkeep")
        imp.changes = False
        imp.close()
        imp = IJ.getImage()
        IJ.selectWindow(imp.getID())

    dimT = imp.getDimensions()[4]

    composite = imp.duplicate()
    composite.setTitle("Composite")
    composite.setDisplayMode(IJ.COMPOSITE)
    IJ.run(composite, "RGB Color", "frames")

    # On ajoute la barre d'échelle si nécessaire
    if scale_length != 0:
        if pixel_size != 0:
            IJ.run(composite, "Set Scale...", "distance=1 known=" + str(pixel_size) + " unit=µm")
        IJ.run(composite, "Scale Bar...",
               "width=" + str(scale_length) + " height=4 font=6 color=White background=None location=[Lower Right] hide overlay")

    if process_as_film:
        montage_stack = ImageStack(imp.getWidth() * (len(separated_channels) + 1), imp.getHeight())   # un cadre de plus que le nombre de canaux, pour le composite

    for t in range(dimT):
        imp.setT(t + 1)

        # Pour chaque canal, on crée une image noir-et-blanc qu'on ajoute au stack
        stack = ImageStack(imp.getWidth(), imp.getHeight())

        for c in separated_channels:
            imp.setC(c)
            grayImp = ImagePlus("temp", imp.getProcessor())
            grayImp.setDisplayRange(imp.getDisplayRangeMin(), imp.getDisplayRangeMax())
            IJ.run(grayImp, "Grays", "")
            IJ.run(grayImp, "RGB Color", "")
            stack.addSlice(grayImp.getProcessor())
            grayImp.close()

        # On finit par ajouter l'image composite au stack
        composite.setT(t)
        stack.addSlice(composite.getProcessor())
        stack_imp = ImagePlus("stack", stack)

        # On crée le montage final à partir du stack, et on le sauvegarde
        montage = MontageMaker().makeMontage2(stack_imp, stack.getSize(), 1, 1, 1, stack.getSize(), 1, 0, False)
        stack_imp.close()

        if process_as_film:
            montage_stack.addSlice(montage.getProcessor())
            montage.close()

    if process_as_film:
        montage = ImagePlus("Montage", montage_stack)

    # On ajoute la barre d'échelle si nécessaire
    if scale_length != 0:
        if pixel_size != 0:
            IJ.run(montage, "Set Scale...", "distance=1 known=" + str(pixel_size) + " unit=µm")
        IJ.run(montage, "Scale Bar...",
               "width=" + str(
                   scale_length) + " height=4 font=6 color=White background=None location=[Lower Right] hide overlay")

    if process_as_film:
        IJ.run(composite, "AVI... ", "compression=JPEG frame=15 save=" + output_basename + "_composite.avi")
        IJ.run(montage, "AVI... ", "compression=JPEG frame=15 save=" + output_basename + "_montage.avi")
    else:
        IJ.saveAs(composite, "PNG", output_basename + "_composite.png")
        IJ.saveAs(montage, "PNG", output_basename + "_montage.png")

    if process_active_image:
        montage.show()
        composite.show()
    else:
        montage.close()
        composite.close()

    return False


## MAIN ##

threshold_popup_location = None
roi_popup_location = None
mins_and_maxs = {}

separated_channels = [int(i) for i in separated_channels]
process_active_image = src_dir is None

if process_active_image:
    processImage()

else:
    interrupt = False  # Utilisé pour arrêter le script avant sa fin

    for root, dirnames, filenames in os.walk(src_dir.getAbsolutePath()):
        dre = re.compile(r'(\d+)')
        filenames.sort(key=lambda l: [int(s) if s.isdigit() else s.lower() for s in re.split(dre, l)])

        for filename in filenames:
            if any([filename.endswith(ext) for ext in [".tif", ".tiff"]]):

                if save_location == "Parent folder":
                    output_basename = root
                else:
                    output_basename = os.path.join(root, os.path.splitext(filename)[0])

                if not any([os.path.isfile(output_basename + output_suffix) for output_suffix in
                        ["_montage.png", "_1_montage.png"]]):
                    interrupt = processImage(os.path.join(root, filename), output_basename)
                    if interrupt:
                        break
        if interrupt:
            break