'''
HYPERSTACK GENERATOR
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

TO DO : 
- Prevent exception if no well selected
- Fix persistence for Boolean
'''
#@PrefService prefs
import os, ij
import ij.plugin as plugin
from ij             import IJ
from fiji.util.gui  import GenericDialogPlus
from java.awt.event import ActionListener

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()[2:]: 
			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()[2:]: 
			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.getBoolean("pickChannel",False)
Channel0     = prefs.get("Channels","") 

pickSlice0   = prefs.getBoolean("pickSlice",False)
Slice0       = prefs.get("Slices","")

pickTime0   = prefs.getBoolean("pickTime",False)
Time0       = prefs.get("Times","")

ShowStack0  = prefs.getBoolean("ShowStack",True)
SaveStack0  = prefs.getBoolean("SaveStack",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)

#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 Path to template
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
	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)
	

	# reformat after the inoput 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 through 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 :
					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: # Channel that were 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:
		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 object
	HyperStacker = plugin.HyperStackConverter() 

	for well in WellList: # Loop over well to do stack and projection

		# 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 (fname,Well,Time,Channel,Slice) for one given well ex: only A001 
		#print 'StackSize : ',len(SubList)
		
		# Initialize stack with first image
		ImagePath = os.path.join(inFolder,SubList[0][0])
		Image     = IJ.openImage(ImagePath)
		Stack     = Image.getStack() 

		for item in SubList[1:]:
			filename = item[0]
			ImagePath = os.path.join(inFolder,filename)
			Image = IJ.openImage(ImagePath)
			
			Stack.addSlice(filename,Image.getProcessor()) # 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
		Stack = ij.ImagePlus(well,Stack) # well is the title of this hyperstack - Reminder : 1 hyperstack/well
		#print Stack
		#Stack.show() # Good

		if 	Stack.getNSlices()>1 : 
			# Convert to Hyperstack
			HyperStack = HyperStacker.toHyperStack(Stack,len(ChannelList), len(SliceList), len(TimeList),"xyczt","grayscale") # convert PROPERLY ORDERED Stack to Hyperstack
		
		else: 
			# Keep as single image
			HyperStack = Stack

			
		if ShowStack : 
			HyperStack.show()

		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)
			
	IJ.log('Done')