/*
 * Align multiple images in a folder
 * saves files in the same folder structures as the input directory
 */

#@ String(value="Batch alignment of images in a folder using StackReg.", visibility="MESSAGE") hint
#@ File (label = "Select input directory", style = "directory") input
#@ File (label = "Select output directory", style = "directory") output
#@ String(value="If you only have one channel, leave the second channel value as is", visibility="MESSAGE") hint2
#@ Integer(label = "Channel 1 (for alignment)", value=1) calcium_channel
#@ Integer(label = "Channel 2 (if none, leave it as is)", value=2) marker_channel
//String(value="Make sure output directory is not within the input directory", visibility="MESSAGE") hint
#@ String (label = "File extension (.czi,.tif,.lif)", value = ".czi") suffix
//setBatchMode(true);

var fs = File.separator;
processFolder(input,output);

exit("Alignment Completed");

function processFolder(input,output) 
{
	list = getFileList(input);
	list = Array.sort(list);
	Array.print(list);
	for (i = 0; i < list.length; i++) 
	{	
		//make sure the output folder is ignored or if any image files within may be opened/aligned
		if(toLowerCase(File.getNameWithoutExtension(list[i]))!=toLowerCase(File.getNameWithoutExtension(output)))
		{
			//if the subdir is a directory, then create a similar output diretory and recursively process folder again
			if(File.isDirectory(input + File.separator + list[i]))
			{
				output_dir=output+ File.separator +list[i];
				if(!File.exists(output_dir)) File.makeDirectory(output_dir);
				print("\nOutput Directory is: "+output_dir);
				processFolder(input + File.separator + list[i]);
			}

			//if it is an image, then process image	
			if(endsWith(list[i], suffix))
			{
				showProgress(i, list.length);
				output_dir=output+fs;
				if(!File.exists(output_dir)) File.makeDirectory(output_dir);
				align_image(input, output_dir, list[i]);
					//print("Processing file");	
			}
		}	
	}
}


//function to align channel 1 for image stacks
function align_image(input, output_dir, file) 
{
	print("Processing: " + input + fs+ file);
	orig_img_path=input + fs+file;
	
	//use bioformats if czi or lif, otherwise use open()
	if(endsWith(orig_img_path, ".czi")|| endsWith(orig_img_path, ".lif")) run("Bio-Formats", "open=["+orig_img_path+"] color_mode=Composite rois_import=[ROI manager] view=Hyperstack stack_order=XYCZT");
	else { open(orig_img_path);}
	stack=getTitle();
	Stack.getDimensions(width, height, channels, slices, frames);
	meta=getMetadata("Info");
	//print(a);
	getVoxelSize(vox_width, vox_height, vox_depth, vox_unit);
	frame_interval=Stack.getFrameInterval();
	fps=1/frame_interval;
	Stack.getUnits(X, Y, Z, Time, Value);

	if(slices>frames)
	{
		print("Swapping frames and slices");
		run("Properties...", "slices="+frames+" frames="+slices);
		selectWindow(stack);
		Stack.getDimensions(width, height, channels, slices, frames);
		print("No of frames are: "+frames" and no of slices "+slices);
	}

	
	if(channels>1)
	{
		//print(calcium_channel,marker_channel);
		selectWindow(stack);
		run("Remove Overlay");
		run("Select None");
		run("Duplicate...", "title=calcium duplicate channels="+calcium_channel);
		calcium_stack=getTitle();
		selectWindow(stack);
		run("Duplicate...", "title=marker duplicate channels="+marker_channel);
		marker_stack=getTitle();
		selectWindow(marker_stack);
		run("Z Project...", "stop=10 projection=[Max Intensity]");
		max_marker=getTitle();
		selectWindow(max_marker);
		saveAs("tiff", output_dir+"Marker2_"+file);
		close();
		close(marker_stack);
		close(stack);
	}
	else 
	{
		calcium_stack=stack;
	}


	print("Aligning calcium stack");
	selectWindow(calcium_stack);
	run("Select None");
	run("Remove Overlay");
	//Stack.setSlice(1);
	Stack.setFrame(1);
	//Registration using StackReg
	run("StackReg", "transformation=[Rigid Body]");
	wait(10);
	selectWindow(calcium_stack);
	setMetadata("Info", meta);
	setVoxelSize(vox_width, vox_height, vox_depth, vox_unit);
	Stack.setFrameRate(fps);
	Stack.setFrameInterval(frame_interval);
	Stack.setUnits(X, Y, Z, Time, Value);
	print("Saving to: " + output_dir);	
	saveAs("tiff", output_dir+"Aligned_"+file);
	close("*");
	run("Collect Garbage");
}

//setBatchMode(false);