# @File(label="Image directory", style="directory") inFolder
# @String(label="Image File Extension", required=false, value=".tif") image_extension
# @String(label="Channel", required=false, value="CO2", choices={"CO2","CO6"}) channel
# @Boolean(label="Show stack ?", value="True") ShowStack
# @Boolean(label="Save Stack ?", description="Stacks will be saved in a 'Stack' folder in the image folder",value = "True") SaveStack
# @String(label="Project stack ?",value="No",choices={"No","avg","median","sd","min","max","sum"}) Project
# @Boolean(label="Save Projection ?",description="Projections will be saved in a 'Projected' folder in the image folder", value = "True") SaveProj
'''
TIME-STACK GENERATOR
Read all tif images from an experiment inFolder that contains single channel tif images for the different wells, channel and slices
The script will automatically reads the well index from the filename (using index in the string - FOLLOWING ACQUIFER-IM04 NAMING CONVENTION ex: -A001--...) to generate the well's stack for the required channel

This current implentation only supports timelapse for each well WITHOUT additional Z-stack

TO DO :
1) Sort images according to timepoints in case the files were not properly sorted. Use tuples (ImagePath, Well, Time) and the sorted function with tuples or use pandas dataframe
2) Improve to take into account Z-Stack or timelapse, currently the script will merge all timepoints to the same stack ie merging several z-stacks together
For sake of simplicity no hyperstack generated yet (maybe later, first put all images ordered by channel slice... then use "stack to hyperstack" see FIJI) 
'''
import os, ij
import ij.plugin as plugin

# Find image files for the given channel
inFolder = str(inFolder) # convert filename object to string object

FileList = []
WellList = []
TimeList = []
for fname in os.listdir(inFolder):
	if fname.endswith(image_extension) and channel in fname:
		FileList.append(fname)
		WellList.append(fname[1:5])
		TimeList.append(int(fname[15:18]))

if len(FileList) < 1:
    raise Exception("No image files found in %s. Check if filenames match IM04 naming convention and contains channel information (COx)" % inFolder)

# Get list of unique well id
WellList = list(set(WellList))
WellList.sort()
TimeStep = max(TimeList)

print 'List of wells', WellList
print 'Number of files : ', len(FileList)
print 'Number of timepoints : ', TimeStep

# Open first image to get the image size (needed to setup the stack) or find another way to do the stack
ImagePath = os.path.join(inFolder,FileList[0])
Image     = ij.IJ.openImage(ImagePath)
Width  = Image.width
Height = Image.height

Projection = ij.ImageStack(Width,Height) # Stack that will contain projected images

for well in WellList: # Loop over well to do stack and projection
	SubFileList = [filename for filename in FileList if well in filename] # isolate pictures composing one stack
	print 'StackSize : ',len(SubFileList)
	
	Stack = ij.ImageStack(Width,Height)
	
	for filename in SubFileList : # Loop over stack slices

		# Add an additional looping here to check for timepoints, if not all the same timepoints then should generate stacks for each of those timepoints
		# time is the T field not the TM (temperature)
		# Watch out the time is precise to second or less therefore for a given timepoint the slice do not have the same time tag, but the order of magnitude of the time difference will of course be larger 
		# between timepoints compared to between slices
		# Alternative solution : ask the user for the number n of Z-stack slices so that we pack the n closest timepoints
		
		ImagePath = os.path.join(inFolder,filename)
		Image = ij.IJ.openImage(ImagePath)
				
		Stack.addSlice(Image.getProcessor()) # conversion to ImageProcessor object to append to stack

	Stack = ij.ImagePlus(well,Stack) # reconverting back to an ImagePlus object to be able to show and save it

	if ShowStack : 
		Stack.show()

	if SaveStack :
		outFolder = os.path.join(inFolder,'Time-Stack')
		if not os.path.exists(outFolder): os.makedirs(outFolder) # create folder if it does not exist (for the first image to save basicly)
		
		outPath = os.path.join(outFolder,well+'.tif')
		ij.IJ.save(Stack,outPath)

	if Project != "No" : # if different from "No" Projection
		Projector = plugin.ZProjector()			 # create projector object
		outImage  = Projector.run(Stack,Project) # project using method define in input

		Projection.addSlice(outImage.getProcessor())
		
		if SaveProj :
			outFolder = os.path.join(inFolder,'Time-Projected')
			if not os.path.exists(outFolder): os.makedirs(outFolder) # create folder if it does not exist (for the first image to save basicly)
		
			outPath = os.path.join(outFolder,well+'.tif')
			ij.IJ.save(Stack,outPath)

# Once the looping is done, show the projected images in a stack
if Project != "No" : 
	Projection = ij.ImagePlus("Time-Projected",Projection)
	Projection.show()