#@ File (label = "Path of model image (let blank to use front-most image as model) ?", required=false) model_image_path
#@ File (label = "Directory to apply LUTs (let blank to apply LUTs on every open image except model)", style = "directory", required=false) src_dir
#@ String (label = "File suffix") suffix

import os
from ij import IJ, WindowManager
from ij.plugin import Duplicator


use_file_image = model_image_path is not None
process_folder = src_dir is not None

## Retrieving the LUTs

model_imp = IJ.openImage(model_image_path) if use_file_image else IJ.getImage()
luts = model_imp.getLuts()
if use_file_image:
    model_imp.close()

## Applying the LUTs

def apply_luts(imp, imp_luts):
    dimC = imp.getDimensions()[2]
    if dimC == len(imp_luts):
        imp.setDisplayMode(IJ.COMPOSITE)
        if dimC > 1:
            imp.setLuts(imp_luts)
        else:
            imp.setLut(imp_luts[0])
        return True
    return False

if process_folder:
    filepaths = []
    for root, dirnames, filenames in os.walk(src_dir):
        for filename in filenames:
            filenames.sort()
            if filename.endswith(suffix):
                filepaths += [os.path.join(root, filename)]

    for filepath in filepaths:
        imp = IJ.openImage(filepath)
        applied = apply_luts(imp, luts)
        if applied:
            IJ.save(imp, filepath)
        imp.close()

else:
    for title in WindowManager.getImageTitles():
        imp = WindowManager.getImage(title)
        apply_luts(imp, luts)