# @File(label = "Input directory", style = "directory") srcFile
# @File(label = "Output directory", style = "directory") dstFile
# @String(label = "File extension", value=".nd2") ext
# @String(label = "File name contains", value = "") containString
# @Integer(label = "Which file number start from?", value = 0) startNum
# @boolean(label = "Keep directory structure when saving", value = true) keepDirectories
# @boolean(label = "Include subfolders?", value = false) subDirectories
# @boolean(label = "Align stack using StackReg?", value = true) align

### This Jython (Python for Java) batch script for Fiji (ImageJ) separates a multidimensional image
### into separate files for each channel and timestep and aligns all the frames to the first one
### using StackReg plugin. It saves the aligned frames in a defined directory together with
### timing information for each frames (that it tries to get from the initial file or separate csv file with the same name)

## The input variables are:
## srcFile: select the directory containing the files that should be processed
## dstFile: select the directory where the final files should be saved
## ext: select the extension of the files that should be processed
## containString: if the directory contains other files, enter the common part of the filenames that should be processed
## startNum: only if filenames end with a number: select the number the script should start from
## align: choose whether you need the images to be aligned with StackReg (true by default)

import os
import csv
from ij import IJ, ImagePlus, WindowManager
from loci.plugins import BF
from loci.formats import ImageReader
from loci.formats import MetadataTools
from ome.xml.meta import OMEXMLMetadata
from xml.etree import ElementTree as ETree
from java.io import File


def run():
    srcDir = srcFile.getAbsolutePath()
    dstDir = (srcFile.getAbsolutePath()
    if str(dstFile) == 'same' else dstFile.getAbsolutePath())

    for (root, directories, filenames) in os.walk(srcDir):
        filenames.sort()
        for filename in filenames:
            
            # Check for file extension
            if not filename.endswith(ext):
                continue
                
            # Check for file name pattern if necessary
            if containString not in filename:
                continue
            
            # Check whether the files in subdirectories should be processed as well
            if not subDirectories:
                if root != srcDir:
                    continue
                    
            # Check whether the filenumber is bigger than 'start from' number
            if startNum != 0 and int(filename[-6:-4]) < startNum:
                continue
            process(srcDir, dstDir, root, filename, keepDirectories)

            
            
def process(
        srcDir,
        dstDir,
        currentDir,
        fileName,
        keepDirectories,
):
    print 'Processing:'
    
    # Opening the image as virtual stack with splitted channels using Bioformats importer
    print 'Opening image file', fileName
    path = os.path.join(currentDir, fileName)
    imp = IJ.run('Bio-Formats', 'open=[' + path +'] color_mode=Default rois_import=[ROI manager] split_channels view=Hyperstack stack_order=XYCZT use_virtual_stack')
    
    # Populating metadata
    file = File(currentDir, fileName)
    r = ImageReader()
    meta = MetadataTools.createOMEXMLMetadata()
    r.setMetadataStore(meta)
    r.setId(file.getAbsolutePath())
    
    # Trying to get DeltaTime for each frame from metadata
    time = []
    for i in range(0, meta.getPlaneCount(0), 2):
        time.append(meta.getPlaneDeltaT(0, i).value())
        
    # if DeltaTime is not found in metadata, trying to load an external csv file with the same name
    # and directory. csv file should contain a single column with timestamps.
    try:
        if time == []:
            print "Can't get timestamps from metadata, trying an external csv file"
            with open(os.path.join(currentDir, fileName[0:-4] + '.csv'), 'rb') as t:
                reader = csv.reader(t)
                time = [row[0] for row in reader]
                t.close()
    except:
        print 'External file with timestamps is not found'
    
    # converting timestamps to seconds if they are in milliseconds 
    try:
        if float(time[0]) > 60:
            time = [x / 1000 for x in time]
    except:
        pass
    time = [[str(x)] for x in time]

    open_images_titles = WindowManager.getImageTitles()
    # iterates over open windows (channels)
    for image_name in open_images_titles:  
        # forms output name
        ch_pos = image_name.find('C=')
        if ch_pos == -1:
            output_name = image_name
        else:
            output_name = image_name[:ch_pos - 7] + '_c' + image_name[ch_pos + 2:]

        # duplicates completely in memory and shows new window
        image = WindowManager.getImage(image_name)
        imageD = image.duplicate()
        imageD.show()
        
        # fix the image to a 12-bit range
        imageD.setDefault16bitRange(12)
        IJ.setMinAndMax(imageD, 0, 4095)
        
        # aligning the images with StackReg plugin
        if align:
            IJ.run(imageD, 'StackReg ', 'transformation=[Translation]')
        
        # forms output directory
        saveDir = (currentDir.replace(srcDir, dstDir) + '//' + image_name[:ch_pos - 7] if keepDirectories else dstDir + '//' + image_name[:ch_pos - 7])
        if not os.path.exists(saveDir):
            print 'Creating folder ', saveDir
            os.makedirs(saveDir)
        print 'Saving to', saveDir
        
        # saves a channel as tiff-sequence: one image for each timestep
        IJ.run(imageD, 'Image Sequence... ', 'format=TIFF name=[' + output_name + '_t] digits=3 save=[' + saveDir + '\\' + output_name + '.tif]')
        print 'File saved as', saveDir + '\\' + output_name + '.tif]'
        
        # closes windows
        imageD.changes = False  # avoid "Save changes?" dialog
        imageD.close()  
        image.close()
    
    # saves timestamps in a separate csv file in the output directory
    with open(os.path.join(saveDir, 'DeltaTime.csv'), 'wb') as f:
        writer = csv.writer(f, dialect='excel', delimiter=',')
        writer.writerows(time)
        f.close()
        
    # saves all metadata in a separate xml file in the output directory
    with open(os.path.join(saveDir, 'Metadata.xml'), 'wb') as f:
        f.write(meta.dumpXML().encode('utf-8'))
        f.close()

run()