'''
HYPERSTACK GENERATOR in Fiji 
Read all tif images from an experiment inFolder that contains tif images for the different wells, channel, slices and timepoints 
The script will automatically reads the field from the filename (using index in the string - FOLLOWING ACQUIFER-IM04 NAMING CONVENTION ex: -A001--...) to generate the well's stack
 
NB : issue if all the images in the folder dont have the same size 
 
TO DO :  
- Add projections to a stack 
- Make plate montage for each time point and/or Z-slice with slider that update the full plate display  
 
Requirements 
- Stack Focuser https://imagej.nih.gov/ij/plugins/stack-focuser.html 
'''
#@PrefService prefs 
import os, ij 
from ij               import IJ, VirtualStack 
from ij.plugin 		  import HyperStackConverter,ZProjector 
from ij.WindowManager import getCurrentImage, getCurrentWindow 
from fiji.util.gui    import GenericDialogPlus 
from java.awt.event   import ActionListener 
from loci.formats 	  import ChannelSeparator   
 
def dimensionsOf(path):   
  fr = None   
  try:   
    fr = ChannelSeparator()   
    fr.setGroupFiles(False)   
    fr.setId(path)   
    return fr.getSizeX(), fr.getSizeY()   
  except:   
    # Print the error, if any   
    print sys.exc_info()   
  finally:   
    fr.close() 
	 
class ListenerClass_All(ActionListener): # extends action listener 
	'''Class associated to the Select all button'''  
	 
	def actionPerformed(this,event): 
		'''Called when 'Select All' is clicked''' 
		global Win 
		 
		for checkbox in Win.getCheckboxes()[-96:]:  
			checkbox.setState(True)  
 
 
class ListenerClass_None(ActionListener): # extends action listener 
	'''Class associated to the Select none button'''  
	 
	def actionPerformed(this,event): 
		'''Called when 'Select None' is clicked''' 
		global Win 
		 
		for checkbox in Win.getCheckboxes()[-96:]:  
			checkbox.setState(False)  
		 
 
ListenerAll  = ListenerClass_All()  # create an instance of the EventListener class 
ListenerNone = ListenerClass_None() # create an instance of the EventListener class 
 
# Recover input from persistence 
Path0        = prefs.get("ImageDirectory","MyPath") 
 
pickChannel0 = prefs.getInt("pickChannel",False) 
Channel0     = prefs.get("Channels","")  
 
pickSlice0   = prefs.getInt("pickSlice",False) 
Slice0       = prefs.get("Slices","") 
 
pickTime0   = prefs.getInt("pickTime",False) 
Time0       = prefs.get("Times","") 
 
ShowStack0  = prefs.getInt("ShowStack",True) 
SaveStack0  = prefs.getInt("SaveStack",False) 
 
ProjMethod0  = prefs.get("ProjMethod","None") 
ShowProj0    = prefs.getInt("ShowProj",False) 
SaveProj0    = prefs.getInt("SaveProj",False) 
 
 
# Input GUI window   
Win = GenericDialogPlus("Hyperstack generator - Parameters") # Title of the window 
 
Win.addDirectoryField("Image directory",Path0) 
Win.addStringField("Image file extension",".tif") 
 
Win.addCheckbox("Use specific channel(s)",pickChannel0) 
Win.addToSameRow() 
Win.addStringField("Channel(s) separated by , (from 1 to 6) : ",Channel0) 
 
Win.addCheckbox("Use specific Z-slice(s)",pickSlice0) 
Win.addToSameRow() 
Win.addStringField("Z-slice(s) separated by , : ",Slice0) 
 
Win.addCheckbox("Use specific timepoint(s)",pickTime0) 
Win.addToSameRow() 
Win.addStringField("Timepoint(s) separated by , : ",Time0) 
 
Win.addCheckbox("Show stack ?",ShowStack0) 
Win.addCheckbox("Save stack ?",SaveStack0) 
 
Win.addChoice("Make Z-projection ?",['None','avg','sd','max','min','sum','median','Extended Focus'],ProjMethod0) 
Win.addToSameRow() 
Win.addCheckbox("Show projection ?",ShowProj0) 
Win.addToSameRow() 
Win.addCheckbox("Save projection ?",SaveProj0) 
 
#Disclaimer = "<html><a href='http://forum.imagej.net/'>This is my link</a></html>" # this will pop up a small window with some documentation 
Disclaimer = "https://github.com/acquifer/ImageInLife/wiki/Stacks"                  # this will directly open the doc page 
Win.addHelp(Disclaimer) 
 
 
# Generate tick box with plate layout to select well 
Wells = [] 
for letter in ['A','B','C','D','E','F','G','H']: 
	for column in range(1,13): 
		Wells.append(letter+'%(column)03d' %{"column":column}) # force max 3 digits and leading 0 
 
Win.addCheckboxGroup(8,12,Wells,[True]*96) 
 
Win.addButton('Select all' ,ListenerAll) # add our event listener to the button. Button clicked -> event -> call actionPerformed(MyEvent) 
Win.addButton('Select none',ListenerNone) 
Win.showDialog() 
 
 
 
## Recover input 
if (Win.wasOKed()):  
	inFolder        = Win.getNextString() 
	image_extension = Win.getNextString() 
	Bool            = [checkbox.getState() for checkbox in Win.getCheckboxes()] 
	 
	i = 0 
	pickChannel     = Bool[i] 
	ChoiceChannel   = Win.getNextString() 
 
	i+=1 
	pickSlice       = Bool[i] 
	ChoiceSlice     = Win.getNextString() 
 
	i+=1 
	pickTime       = Bool[i] 
	ChoiceTime     = Win.getNextString() 
 
	i+=1 
	ShowStack      = Bool[i] 
	 
	i+=1 
	SaveStack      = Bool[i] 
	 
	i+=1 
	ProjMethod 	   = Win.getNextChoice() 
	ShowProj       = Bool[i] 
	 
	i+=1 
	SaveProj 	   = Bool[i]	 
 
	i+=1 
	BoolWell       = Bool[i:]  
	Selection      = {well:choice for well,choice in zip(Wells,BoolWell)} # create a dictionnary associating a well to a boolean . Chosen or not 
	#print Selection 
	#print len(Wells) # OK 96 
 
	 
	# Save input into memory/persistence 
	prefs.put("ImageDirectory",inFolder) 
	 
	prefs.put("pickChannel",pickChannel) 
	prefs.put("Channels",ChoiceChannel) 
 
	prefs.put("pickSlice",pickSlice) 
	prefs.put("Slices",ChoiceSlice) 
 
	prefs.put("pickTime",pickTime) 
	prefs.put("Times",ChoiceTime) 
 
	prefs.put("ShowStack",ShowStack) 
	prefs.put("SaveStack",SaveStack) 
 
	prefs.put("ProjMethod",ProjMethod) 
	prefs.put("ShowProj",ShowProj) 
	prefs.put("SaveProj",SaveProj) 
	 
	# Stop execution if no well selected 
	if BoolWell == [False]*96: 
		IJ.error("No well selected") 
		'''
		message = "No well selected" 
		IJ.log(message) 
		raise Exception(message) 
		'''
 
	# reformat after the input has been saved in memory 
	if ChoiceChannel != "": # prevent bug when field empty and we convert to int 
		#ChoiceChannel = list(map(int,ChoiceChannel.split(","))) 
		ChoiceChannel = eval(ChoiceChannel) 
		if type(ChoiceChannel)==int: # in case it is a single integer turn it to a list 
			ChoiceChannel = [ChoiceChannel] 
			   
		 
	if ChoiceSlice != "":  
		ChoiceSlice = eval(ChoiceSlice) 
		if type(ChoiceSlice)==int: # in case it is a single integer turn it to a list 
			ChoiceSlice = [ChoiceSlice] 
 
	if ChoiceTime != "": 
		ChoiceTime = eval(ChoiceTime) 
		if type(ChoiceTime)==int: # in case it is a single integer turn it to a list 
			ChoiceTime = [ChoiceTime] 
 
 
	# Find image files for the given channel 
	#inFolder = str(inFolder) # convert filename object to string object 
	WellList    = [] # full list of well to know how many iteration for stacking 
	TimeList    = [] 
	ChannelList = [] 
	SliceList   = [] 
	MainList    = [] # contains tuple with (Filename, Well, TimePoint, Channel, Z-Slice) 
 
	for fname in os.listdir(inFolder): # extract metadata from image filename  
		if fname.endswith(image_extension): 
 
			# first extract the well field 
			Well = fname[1:5] 
 
			# if well is selected in the gui, then extract the other metadata 
			if Selection[Well]: 
 
				try : # int conversion throw an error if the argument string is empty, hence the try  
					Time 	= int(fname[15:18]) 
					Channel = int(fname[22:23]) 
					Slice 	= int(fname[27:30])  
				except : 
					IJ.error("Could not get experimental parameters from filename. Check filename pattern (IM04 convention)") 
					#raise Exception("Could not get experimental parameters from filename. Check filename pattern (IM04 convention)") 
				 
				WellList.append(Well) 
 
				 
				## Channel 
				if not pickChannel: # we take all channel 
					ChannelList.append(Channel) 
					 
				elif pickChannel and (Channel in ChoiceChannel): # We take only a set of timepoints 
					ChannelList.append(Channel) 
 
				else: # This particular channel was not selected 
					continue 
 
					 
				 
				## Time 
				if not pickTime: # we take all time points 
					TimeList.append(Time) 
					 
				elif pickTime and (Time in ChoiceTime): # We take only a set of timepoints 
					TimeList.append(Time) 
 
				else: # Timepoints that were not selected 
					continue 
				 
 
				 
				## Z-Slice 
				if not pickSlice: # we take all time points 
					SliceList.append(Slice) 
					 
				elif pickSlice and (Slice in ChoiceSlice): # We take only a set of timepoints 
					SliceList.append(Slice) 
 
				else: # Timepoints that were not selected 
					continue 
 
				 
				# Append to main list only if it managed to pass the previous if (thanks to the continue statement) 
				MainList.append((fname,Well,Time,Channel,Slice)) 
 
			else : # the well was not selected in the GUI, go for the next image filename 
				continue 
 
			 
	if len(MainList) < 1: 
		IJ.error("No image files found in %s satisfying the channel/slice/time selection. Check if extension match and if filenames match IM04 naming convention" % inFolder) 
		#raise Exception("No image files found in %s satisfying the channel/slice/time selection. Check if extension match and if filenames match IM04 naming convention" % inFolder) 
 
	# Sort list to prepare iteration 
	WellList = list(set(WellList)) # turn to a set object to have unique item once only  
	WellList.sort() 			   # sort should be after set since set may change order  
 
	TimeList = list(set(TimeList))  
	TimeList.sort() 			   
 
	ChannelList = list(set(ChannelList))  
	ChannelList.sort() 
 
	SliceList = list(set(SliceList))  
	SliceList.sort()  
 
	# Verification 
	print 'Well  : ', WellList 
	print 'Time  : ', TimeList 
	print 'ChannelIdx  : ', ChannelList 
	print 'Z-Slice : ' , SliceList 
 
 
	MainList.sort(key = lambda field:(field[1], field[2], field[4], -field[3]) ) # sort by Well, Time, Z-slice then Channel (- to have C06 ie BF first) since the hyperstack is created following the default "xyczt" order 
	'''
	for i in MainList : 
		print i 
	'''
	 
	# Initialise HyperStackConverter and Projector object 
	HyperStacker = HyperStackConverter()  
	Proj         = ZProjector() 
	 
	## Loop over wells to do stack and projection 
	for well in WellList: # - TO DO: Add another level in MainList : one per well 
 
		# Gather images for a given well, and put them into a sublist 
		IJ.log('Process well : '+well+'...') 
		SubList = [item for item in MainList if item[1]==well] # Sublist containing (fname1,Well,Time1,ChannelX,Slice),(fname2,Well,Time2,ChannelX,Slice)  for one given well ex: only A001  
		#print 'StackSize : ',len(SubList) 
		 
		# Get image dimensions of the first image in the list to initialise VStack 
		width, height = dimensionsOf( os.path.join(inFolder,SubList[0][0]) ) 
		 
		# Initialise Virtual Stack for this well 
		VStack = VirtualStack(width, height, None, inFolder)   
		 
		# Append slices in the VStack 
		for item in SubList: 
			filename = item[0] 
			VStack.addSlice(filename) # conversion to ImageProcessor object to append to stack 
	 
			# Once we have finished appending slices, we reconvert back to an ImagePlus object to be able to show and save it 
			ImpStack = ij.ImagePlus(well, VStack) # well is the title of this hyperstack - Reminder : 1 hyperstack/well 
			#print ImpStack
			#ImpStack.show() # Good 
			 
		 
		## Make a hyperstack out of the ImagePlus virtual stack
		if ImpStack.getStackSize()==1: 
			HyperStack = ImpStack # issue when a single slice in the Hyperstack
		else:
			HyperStack = HyperStacker.toHyperStack(ImpStack,len(ChannelList), len(SliceList), len(TimeList),"xyczt","grayscale") # convert PROPERLY ORDERED Stack to Hyperstack 
		 
		 
		## Display resulting stack 
		if ShowStack :  
			HyperStack.show() 
		 
		 
		## Save resulting stack 
		if SaveStack : 
			outFolder = os.path.join(inFolder,'HyperStack') 
			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.save(HyperStack,outPath)		 
		 
		 
		## PROJECTION 
		if ProjMethod != 'None': 
	 
			# Report 
			message = 'Do projection' 
			#print message 
			IJ.log(message) 
			 
			# Do projection 
			if ProjMethod == 'Extended Focus': 
				IJ.run(HyperStack, "Stack Focuser ", "enter=3") # return void but open a new image window 
				Projected = getCurrentImage().duplicate() 		# duplicate otherwise lost once the window is closed 
				 
				if not ShowProj: 
					getCurrentWindow().close() 
 
			 
			else: # other classical projection method 
				Projected = Proj.run(HyperStack,ProjMethod) 
			 
				# Display projection (automatic with the stack focuser) 
				if ShowProj: 
					Projected.show() 
 
			 
			# Save projection 
			if SaveProj: 
				outFolder = os.path.join(inFolder,'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.save(Projected,outPath)