# @File(label="Image directory", style="directory") inFolder
# @String(label="Channel", required=false, value="CO2", choices={"CO1","CO2","CO3","CO4","CO5","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="Show Projection ?", value = "True") ShowProj
# @Boolean(label="Save Projection ?",description="Projections will be saved in a 'Projected' folder in the image folder", value = "True") SaveProj
'''
Z-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 support only Z-stack WITHOUT additionnal timelapse for eac hwell

TO DO :
1) Sort images according to Z slice in case the files were not properly sorted. Use tuples (ImagePath, Well, Z) and the sorted function with tuple 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 = []
for fname in os.listdir(inFolder):
	filepath = os.path.join(inFolder,fname)
	
	if os.path.isfile(filepath) and channel in fname:
		FileList.append(fname)
		WellList.append(fname[1:5])

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)

# 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

# Get list of unique well id
WellList = list(set(WellList))
WellList.sort()
#print 'List of wells', WellList
#print 'Number of files : ', len(FileList)

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

for well in WellList: # Loop over well position (A001,A002...) 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) # Stack for one well
	
	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(filename,Image.getProcessor()) # conversion to ImageProcessor object to append to stack

	
	Stack = ij.ImagePlus(well,Stack) # reconverting Stack of one well to an ImagePlus object to be able to show and save it
	
	if ShowStack : 
		Stack.show()

	if SaveStack :
		outFolder = os.path.join(inFolder,'Z-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())
		


# Once the looping is done, show/save the stack of projected images
if Project != "No" : 
	Projection = ij.ImagePlus("Z-Projected",Projection) # convert back to an ImagePlus 

	if ShowProj :
		
		Projection.show()
	
	if SaveProj :
		# Save projections as a stack
		outFolder = os.path.join(inFolder,'Z-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,'Z-Projected.tif')
		ij.IJ.save(Projection,outPath) # overwrite current image ?