// Correct_SoRa macro by Christophe Leterrier
// 13/03/2025
//
// 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
// Update 16/06/2025: replace ONE Microscopy NanoJ Core version by Fast4DReg version

macro "Correct_SoRa" {

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

// Save Settings
	saveSettings();

// Default variables
	FULL_WIDTH = 2304; // full snesor width in pixels

// Default values for the Options Panel
	OBJ_DEF = "60X"
	CONF_DEF = "SoRa";
	MODE_DEF = "CZ";
	CROP_DEF = "Full";
	CHAN_DEF = "0,1,2";
	SPLIT_DEF = false;
//	OPT_DEF = false;
	SAVE_DEF = "In a folder next to the source folder";

// Initialize choices variables
	OBJ_ARRAY = newArray("100X", "60X", "40X");
	CONF_ARRAY = newArray("Confocal", "SoRa");
	MODE_ARRAY = newArray("CZ", "ZC");	
	CROP_ARRAY = newArray("Full", "Crop");
	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");
	run("Close All");


//*************** 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 SoRa Log ***");
	print("INPUT_DIR: " + INPUT_DIR);


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

	//Creation of the dialog box
	Dialog.create("Correct SoRa Options");
	Dialog.addChoice("    Objective", OBJ_ARRAY, OBJ_DEF);
	Dialog.addChoice("    Confocal Mode", CONF_ARRAY, CONF_DEF);
	Dialog.addChoice("     Stack mode", MODE_ARRAY, MODE_DEF);
	Dialog.addChoice("     Crop mode", CROP_ARRAY, CROP_DEF);
	Dialog.addMessage("Channels: 0/1/2=640/561/488");
	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
	OBJ = Dialog.getChoice();
	CONF = Dialog.getChoice();
	MODE = Dialog.getChoice();
	CROP = 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 SoRa" + File.separator + "Chroma" + File.separator + OBJ + File.separator + MODE + File.separator + CONF + 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_FULL_PATH = CMASK_DIR + File.separator + "CustomMask.tif";
	if (File.isDirectory(CMASK_DIR) == false) {
			File.makeDirectory(CMASK_DIR);
	}
	save(CMASK_FULL_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();
			FILE_INFO = getMetadata("Info");
			getDimensions(IN_WIDTH, IN_HEIGHT, N_CHANNELS, N_SLICES, N_FRAMES);

			// Only process images that have the declared number of channels
			if (N_CHANNELS == CHAN_ARRAY.length){
				
				// Crop mask if option selected
				if (CROP == "Crop") {
					
					// Get left and top coordinate of image on sensor (from Nikon metadata)
					CornerCoor = getPropValues("left,top");
					print("ROI detected: top= " + CornerCoor[0] + " left= " + CornerCoor[1]);
					
					// Only generate custom cropped mask if image is cropped
					if (CornerCoor[0]!=0 || CornerCoor[1]!=0) {				
						// Open custom mask, crop it to the same area
						open(CMASK_FULL_PATH);
						CMASK_FULL_ID = getImageID();
						makeRectangle(CornerCoor[0], CornerCoor[1], IN_WIDTH, IN_HEIGHT);
						run("Duplicate...", "duplicate");
						rename("Crop1");
						selectImage(CMASK_FULL_ID);
						makeRectangle(CornerCoor[0]+FULL_WIDTH, CornerCoor[1], IN_WIDTH, IN_HEIGHT);
						run("Duplicate...", "duplicate");
						rename("Crop2");
						run("Combine...", "stack1=Crop1 stack2=Crop2");
						CMASK_CROP_ID = getImageID();
						
						// Save cropped mask
						selectImage(CMASK_FULL_ID);
						close();
						selectImage(CMASK_CROP_ID);
						CMASK_CROP_PATH = CMASK_DIR + File.separator + FILE_SHORTNAME + "_CustomMask.tif";
						save(CMASK_CROP_PATH);
						close();
						
						// Switch Mask path to cropped mask path
						CMASK_PATH = CMASK_CROP_PATH;
					}
					
					// Else use full custom mask
					else CMASK_PATH = CMASK_FULL_PATH;
				}
				
				// If no crop option, use full custom mask
				else CMASK_PATH = CMASK_FULL_PATH;
					
				// Case of z stacks (multi-channel 3D)
				if (N_SLICES > 1) {
					print("Multi-channel Z-stack detected");
					// Duplicate input hyperstack (to generate output hyperstack)
					selectImage(FILE_ID);
					rename("Input stack");
					run("Duplicate...", "duplicate");
					rename("Aligned stack");
					OUT_ID = getImageID();
						
					for (z=0; z < N_SLICES; z++) {						
						// Duplicate single slice and convert into regular stack
						selectImage(FILE_ID);
						run("Duplicate...", "duplicate slices=" + z+1);
						run("Hyperstack to Stack");
						UNALIGNED_ID = getImageID();			
						
						// Run aberration correction using custom mask, get output title and ID
						run("F4DR Correct Channel Misalignment", "open=[" + CMASK_PATH + "]");
						// run("Register Channels - Apply", "open=[" + CMASK_PATH + "]");
						ALIGNED_TITLE = getTitle();
						ALIGNED_ID = getImageID();
						
						// Copy/paste each channel of the output into the output hyperstack
						for (j=0; j < N_CHANNELS; j++) {
							selectImage(ALIGNED_ID);
							setSlice(j+1);
							run("Select All");
							run("Copy");
							
							selectImage(OUT_ID);
							Stack.setPosition(j+1, z+1, 1);
							run("Paste");		
						}
						
						selectImage(UNALIGNED_ID);
						close();
						selectImage(ALIGNED_ID);
						close();
					}
					
					selectImage(OUT_ID);
					Stack.setPosition(1,1,1);
				}
				
				// Case of time lapse (multi-channel 2D+t)
				else if (N_FRAMES > 1) {
					print("Multi-channel timelapse detected");
					// Duplicate input hyperstack (to generate output hyperstack)
					selectImage(FILE_ID);
					run("Duplicate...", "duplicate");
					OUT_ID = getImageID();
					
					for (t=0; t < N_FRAMES; t++) {	
						// Duplicate single slice and convert into regular stack
						selectImage(FILE_ID);
						run("Duplicate...", "duplicate frames=" + t+1);
						run("Hyperstack to Stack");
						UNALIGNED_ID = getImageID();
						
						// Run aberration correction using custom mask, get output title and ID
						run("Hyperstack to Stack");
						// run("Register Channels - Apply", "open=[" + CMASK_PATH + "]");
						run("F4DR Correct Channel Misalignment", "open=[" + CMASK_PATH + "]");
						ALIGNED_TITLE = getTitle();
						ALIGNED_ID = getImageID();
						
						// Copy/paste each channel of the output into the output hyperstack
						for (j=0; j < N_CHANNELS; j++) {
							selectImage(ALIGNED_ID);
							setSlice(j+1);
							run("Select All");
							run("Copy");
							
							selectImage(OUT_ID);
							Stack.setPosition(j+1, t+1, 1);
							run("Paste");		
						}
						
						// Needed because the Hyperstack to stack does not create a new stack for 2D multichannel hyperstack
						// FILE_ID = UNALIGNED_ID;
						selectImage(UNALIGNED_ID);
						close();
						selectImage(ALIGNED_ID);
						close();
					}
				
					selectImage(OUT_ID);
					Stack.setPosition(1,1,1);
				
				}
				
				// Case of single images (multi-channel 2D)
				else {
					
					print("Multi-channel single plane detected");
					// Duplicate input stack (to generate output stack)
					selectImage(FILE_ID);
					run("Duplicate...", "duplicate");
					OUT_ID = getImageID();
					
					selectImage(FILE_ID);
					run("Hyperstack to Stack");
					UNALIGNED_ID = getImageID();
					
					// Run aberration correction using custom mask, get output title and ID
					//run("Register Channels - Apply", "open=[" + CMASK_PATH + "]");
					run("F4DR Correct Channel Misalignment", "open=[" + CMASK_PATH + "]");
					ALIGNED_TITLE = getTitle();
					ALIGNED_ID = getImageID();
					
					// Copy/paste each channel of the output into the output hyperstack
					for (j=0; j < N_CHANNELS; j++) {
						selectImage(ALIGNED_ID);
						setSlice(j+1);
						run("Select All");
						run("Copy");
						
						selectImage(OUT_ID);
						Stack.setPosition(j+1, 1, 1);
						run("Paste");
						resetMinAndMax;
					}
					
					// Needed because the Hyperstack to stack does not create a new stack for 2D multichannel hyperstack
					FILE_ID = UNALIGNED_ID;
					// selectImage(UNALIGNED_ID);
					// close();
					selectImage(ALIGNED_ID);
					close();
				
					selectImage(OUT_ID);
					Stack.setPosition(1,1,1);
				}
				
				// 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);
	
						// set Metadata
						run("Select None");
						setMetadata("Info", FILE_INFO);
	
						// 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);
					
					// set Metadata
					run("Select None");
					setMetadata("Info", FILE_INFO);
	
					// 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 SoRa end ***");
	showStatus("Correct SoRa 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;
}

function getPropValues(keys) {
	// get values of image properties defined by a "keys" string with property names separated by commas
	// return an array of values correspoonding to each key

	// define delimiters between a property and its value, with a parenthesed version for the split 
	del1 = " = ";
	del2 = "; ";

	// get properties
	Info = Property.getInfo();
	// split by lines
	InfoArray = split(Info, "\n");
	pL = InfoArray.length;

	PropArray = newArray(pL); // array of property names
	ValArray = newArray(pL); // array of property values

	// fill the property names and values arrays
	for (l = 0; l < pL; l++) {
		
		// split each line according to delimiters defined above (del1, del2) or assign NaN if no delimiter detected
		if (indexOf(InfoArray[l], del1) > -1) LineSplit = split(InfoArray[l], "(" + del1 + ")"); // the split string is parenthesed to work as regular expression
			else if (indexOf(InfoArray[l], del2) > -1) LineSplit = split(InfoArray[l], "(" + del2 + ")");
			else LineSplit = newArray(NaN, NaN);
	
		PropArray[l] = LineSplit[0];
		ValArray[l] = LineSplit[1];

	}

	// split keys (asked property names)
	KeyArray = split(keys, ",");
	kL = KeyArray.length;
	KeyValArray = newArray(kL);

	// for each key of KeyArray, look in the PropArray property names array, and when name is found, assign corresponding value from ValArray to the corresponding KeyValArray slot
	for (k = 0; k < kL; k++) {
		flag = 0;
		for (l = 0; l < pL; l++) {
			if (PropArray[l] == KeyArray [k]) {
				flag = 1;
				KeyValArray[k] = ValArray[l];
			}
		}
		if (flag == 0) KeyValArray[k] = NaN; // if key not found, assign NaN for that slot
	}

	// return the key values array
	return KeyValArray;
}