// Correct_SIM macro by Christophe Leterrier
// 06/07/2023
//
// Takes a folder of multi-channel tif and performs:
// - Chromatic shift correction from a NanoJ translation mask obtained on a set of reference images (beads)
// - Splitting of channels into individual tif images
// The extracted images are located in a folder defined in the menu.
//
// Macro uses NanoJ-Core to register channels

macro "Correct_SIM" {

//*************** Initialization ***************

// Save Settings
	saveSettings();

// Default values for the Options Panel
	MODE_DEF = "TIRF-SIM";
	REC_DEF = "HiFi-SIM";
	CHAN_DEF = "0,1,2,3";
	SPLIT_DEF = false;
	OPT_DEF = false;
	SAVE_DEF = "In a folder next to the source folder";
	run("Close All");

// Initialize choices variables
	MODE_ARRAY = newArray("TIRF-SIM", "2D-SIM", "3D-SIM");
	REC_ARRAY = newArray("NSIM", "HiFi-SIM", "Wiener-SIM", "Widefield");	
	SAVE_ARRAY = newArray("In the source folder", "In a subfolder of the source folder", "In a folder next to the source folder", "In a custom folder");



//*************** Dialog 1 : get the input images folder path ***************

	INPUT_DIR = getDirectory("Select a source folder with multi-channel tif images");

	print("\n\n\n*** Correct SIM Log ***");
	print("INPUT_DIR: " + INPUT_DIR);


//*************** Dialog 2 : options ***************

	//Creation of the dialog box
	Dialog.create("Correct SIM Options");
	Dialog.addChoice("     SIM mode", MODE_ARRAY, MODE_DEF);
	Dialog.addChoice("Reconstruction", REC_ARRAY, REC_DEF);
	Dialog.addMessage("Channels: 0/1/2/3=640/561/488/405");
	Dialog.addString("     Channels", CHAN_DEF, 15);
	Dialog.addCheckbox("Split Channels", SPLIT_DEF);
	Dialog.addCheckbox("Enhance Contrast", OPT_DEF);
	Dialog.addChoice("Save Images", SAVE_ARRAY, SAVE_DEF);
	Dialog.show();

	// Feeding variables from dialog choices
	MODE = Dialog.getChoice();
	REC = Dialog.getChoice();
	CHAN = Dialog.getString();
	SPLIT_CH = Dialog.getCheckbox();
	OPT = Dialog.getCheckbox();
	SAVE_TYPE = Dialog.getChoice();

	CHAN_ARRAY = split(CHAN, ",");
	for ( i = 0; i < CHAN_ARRAY.length; i++) {
		CHAN_ARRAY[i] = parseInt(CHAN_ARRAY[i]);
	}

// Set the path to the scope correction images, name the output channel "Corrected" or "Uncorrected"
	OUTSTRING = "shift";
	SCOPE_CORR_DIR = findCorrDir();
	// If not found asks for the folder
	if (File.isDirectory(SCOPE_CORR_DIR) == false) {
		SCOPE_CORR_DIR = getDirectory("Select the correction images folder");
	}
	print("Scope Correction folder path: " + SCOPE_CORR_DIR);

// Add "Split" in the output folder name if the split channel option is chosen
	if (SPLIT_CH == true) OUTSTRING = OUTSTRING + " split";

	setBatchMode(true);

//*************** Prepare Processing (get names, open images, make output folder) ***************

	// Get the path to the correction images and open them
	print("Chroma Correction: ON");
	CHROMA_PATH = SCOPE_CORR_DIR + "Ref SIM" + File.separator + "Chroma" + File.separator + MODE + File.separator + REC + File.separator;
	print("Chromatic correction folder path: " + CHROMA_PATH);

	// Get all file names
	ALL_NAMES = getFileList(INPUT_DIR);
	Array.sort(ALL_NAMES);
	N_LENGTH = ALL_NAMES.length;
	ALL_EXT = newArray(N_LENGTH);
	// Create extensions array
	for (i = 0; i < N_LENGTH; i++) {
	//	print(ALL_NAMES[i]);
		ALL_NAMES_PARTS = getFileExtension(ALL_NAMES[i]);
		ALL_EXT[i] = ALL_NAMES_PARTS[1];
	}

	//Create the output folder
	OUTPUT_DIR = INPUT_DIR;
	NAME_SUFF = "";

	if (SAVE_TYPE == "In the source folder") {
		OUTPUT_DIR=INPUT_DIR;
		NAME_SUFF = "_shift";
	}

	if (SAVE_TYPE == "In a subfolder of the source folder") {
		OUTPUT_DIR = INPUT_DIR + OUTSTRING + File.separator;
		if (File.isDirectory(OUTPUT_DIR) == false) {
			File.makeDirectory(OUTPUT_DIR);

		}
	}

	if (SAVE_TYPE == "In a folder next to the source folder") {
		OUTPUT_DIR = File.getParent(INPUT_DIR);
		OUTPUT_NAME = File.getName(INPUT_DIR);
		OUTPUT_SHORTA = split(OUTPUT_NAME, " ");
		OUTPUT_SHORT = OUTPUT_SHORTA[0];
		OUTPUT_DIR = OUTPUT_DIR + File.separator + OUTPUT_SHORT + " " + OUTSTRING + File.separator;
		if (File.isDirectory(OUTPUT_DIR) == false) {
			File.makeDirectory(OUTPUT_DIR);
		}
	}

	if (SAVE_TYPE == "In a custom folder") {
		OUTPUT_DIR = getDirectory("Choose the save folder");
	}

	OUTPUT_PARENT_DIR = File.getParent(OUTPUT_DIR);

	print("OUTPUT_DIR: " + OUTPUT_DIR);
	//print("OUTPUT_PARENT_DIR: " + OUTPUT_PARENT_DIR);


	//Prepare the translation mask
	// There's no management of the image size (has to be the same as source images)
	
	// Open the reference mask with all channels
	RMASK_PATH = CHROMA_PATH + "TranslationMask.tif";
	open(RMASK_PATH);
	RMASK_ID = getImageID();
	getDimensions(mw, mh, mch, msl, mfr);
	
	// Make a custom mask with the right channels
	newImage("CustomMask.tif", "32-bit black", mw, mh, CHAN_ARRAY.length);
	CMASK_ID = getImageID();
	
	for (j = 0; j<CHAN_ARRAY.length; j++) {
		selectImage(RMASK_ID);
		setSlice(CHAN_ARRAY[j]+1);
		run("Select All");
		run("Copy");
		selectImage(CMASK_ID);
		setSlice(j+1);
		run("Paste");
	}
	
	// Save the custom mask in a folder of the output folder (to avoid a spurious tif in the folder itself)
	CMASK_DIR = OUTPUT_DIR + File.separator + "Translation Mask";
	CMASK_PATH = CMASK_DIR + File.separator + "CustomMask.tif";
	if (File.isDirectory(CMASK_DIR) == false) {
			File.makeDirectory(CMASK_DIR);
	}
	save(CMASK_PATH);
	close();
	
	// close Ref mask
	close();
	


//*************** Process Images ***************

	flag = 0;
	// Loop on all .tif extensions
	for (n = 0; n < N_LENGTH; n++) {
		LOW_EXT = toLowerCase(ALL_EXT[n]);
		if (LOW_EXT == ".tif" || LOW_EXT == ".tiff") {

			flag = flag + 1;
			// Get the file path
			FILE_PATH = INPUT_DIR + ALL_NAMES[n];

			// Store components of the file name
			FILE_NAME = File.getName(FILE_PATH);
			FILE_DIR = File.getParent(FILE_PATH);
			FILE_SEP = getFileExtension(FILE_NAME);
			FILE_SHORTNAME = FILE_SEP[0];
			FILE_EXT = FILE_SEP[1];

			print("");
			print("INPUT_PATH:", FILE_PATH);

			// Open input image
			open(FILE_PATH);
			FILE_TITLE = getTitle();
			FILE_ID = getImageID();
			N_CHANNELS = nSlices;

			// Only process images that have the declared number of channels
			if (N_CHANNELS == CHAN_ARRAY.length){
				
				// Convert input into regular stack
				run("Hyperstack to Stack");
				FILE_ID = getImageID();
				
				// Run aberration correction using custom mask
				//run("Register Channels - Apply", "open=[" + CMASK_PATH + "]");	
				run("F4DR Correct Channel Misalignment", "open=[" + CMASK_PATH + "]");
				
				// Convert output to multi-channel hyperstack
				run("Stack to Hyperstack...", "order=xyczt(default) channels=" + N_CHANNELS + " slices=1 frames=1 display=Grayscale");	
				
				for (j=0; j < N_CHANNELS; j++) {
					setSlice(j+1);
					if (OPT == true) run("Enhance Contrast", "saturated=0.35");
					else resetMinAndMax;
				}
				setSlice(1);
				
				OUT_TITLE = getTitle();
				OUT_ID = getImageID();
				
				// close input
				selectImage(FILE_ID);
				close();				
						
				// Split channels if checked
				if (SPLIT_CH == true) {

					CHANNEL_COUNT = nSlices;
					SLICE_NAMES = newArray(CHANNEL_COUNT);
	
					// Loop on slices to get the images titles after "Stack to Images"
					for(j = 0; j < CHANNEL_COUNT; j++) {
						setSlice(j+1);
						SLICE_NAMES[j] = getInfo("slice.label");
						if (SLICE_NAMES[j] == "") {
							SLICE_NAMES[j] = FILE_SHORTNAME + "-000" + (j+1);
						}
					}
					// Break into single images
					run("Stack to Images");
	
					// Loop on each channel (each opened window)
					for(j = 0; j < CHANNEL_COUNT; j++) {
	
						//Select source image
						selectWindow(SLICE_NAMES[j]);
	
						// rename image according to processing
						NEW_WINDOW_NAME = FILE_SHORTNAME + NAME_SUFF + "-C=" + j;
						rename(NEW_WINDOW_NAME);
	
						// Create output file path and save the output image
						OUTPUT_PATH = OUTPUT_DIR + NEW_WINDOW_NAME + NAME_SUFF + ".tif";
						save(OUTPUT_PATH);
						print("OUTPUT_PATH: " + OUTPUT_PATH);
						close();
						} // end of FOR loop on channels
				}
				
				else {
					// rename image according to processing
					NEW_WINDOW_NAME = FILE_SHORTNAME + NAME_SUFF + ".tif";
					rename(NEW_WINDOW_NAME);
	
					// Create output file path and save the output image
					OUTPUT_PATH = OUTPUT_DIR + NEW_WINDOW_NAME;
					save(OUTPUT_PATH);
					print("OUTPUT_PATH: " + OUTPUT_PATH);
					close();
				}
										
														
			} 	// end of IF loop on number of channels
		}	// end of IF loop on tif extensions
	}	// end of FOR loop on n extensions

// selectWindow("CustomMask.tif");
// close();

//*************** Cleanup and end ***************

	// Restore settings
	restoreSettings();

	setBatchMode("exit and display");

	print("");
	print("*** Correct SIM end ***");
	showStatus("Correct SIM finished");
//	exec("open", OUTPUT_DIR);
}


//*************** Functions ***************

function getFileExtension(Name) {
	nameparts = split(Name, ".");
	shortname = nameparts[0];
	if (nameparts.length > 2) {
		for (k = 1; k < nameparts.length - 1; k++) {
			shortname += "." + nameparts[k];
		}
	}
	extname = "." + nameparts[nameparts.length - 1];
	namearray = newArray(shortname, extname);
	return namearray;
}

function findCorrDir() {
	IJ_PATH = getDirectory("imagej");
	CORR_PATH = IJ_PATH + "scripts" + File.separator + "NeuroCyto" + File.separator + "Scope Correction" + File.separator + "Data" + File.separator;
	return CORR_PATH;
}