// Correct_STORM macro by Christophe Leterrier
// 27-01-2014
// last version 15-04-2021
//
// Takes a folder of multi-channel tif and performs:
// - Chromatic shift correction from a set of reference images
// - Splitting of channels into individual tif images
// The extracted images are located in a folder defined in the menu.
//
// Macro uses TransformJ to register channels, using the matrices calculated by the Get_Chroma_Coefs macro


macro "Correct_STORM" {


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

// Save Settings
	saveSettings();

// Default values for the Options Panel
	MAG_DEF = "100X";
	LENS_DEF = false;
	CAM_DEF = "256x256";
	CHAN_DEF = "0,1,2,3";
	SPLIT_DEF = false;
	SAVE_DEF = "In a folder next to the source folder";

// Initialize choices variables
	MAG_ARRAY = newArray("100X");
	CAM_ARRAY = newArray("512x512", "256x256");
	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 STORM Log ***");
	print("INPUT_DIR: " + INPUT_DIR);


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

	//Creation of the dialog box
	Dialog.create("Correct STORM Options");
	Dialog.addChoice("     Objective", MAG_ARRAY, MAG_DEF);
	Dialog.addChoice("     Field of view", CAM_ARRAY, CAM_DEF);
	Dialog.addMessage("Channels: 0 = NSTORM\n                1/2/3/8 = Epi 647/561/488/405\n                4/5/6/7 = Laser 647/561/488/405\n                9=Trans");
	Dialog.addString("     Channels", CHAN_DEF, 15);
	Dialog.addCheckbox("Split Channels", SPLIT_DEF);
	Dialog.addChoice("Save Images", SAVE_ARRAY, SAVE_DEF);
	Dialog.show();

	// Feeding variables from dialog choices
	MAG = Dialog.getChoice();
	CAM = Dialog.getChoice();
	CHAN = Dialog.getString();
	SPLIT_CH = Dialog.getCheckbox();
	SAVE_TYPE = Dialog.getChoice();

	if (CAM == "512x512") CM = "WiF"; else CM = "ROI";

	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 = "corr";
	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 + "Chroma" + File.separator + MAG + "_" + CM + 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;

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

	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 + " ext " + 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);


//*************** 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);
			//print("FILE_NAME:", FILE_NAME);
			//print("FILE_DIR:", FILE_DIR);
			//print("FILE_EXT:", FILE_EXT);
			//print("FILE_SHORTNAME:", FILE_SHORTNAME);

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

			N_CHANNELS = nSlices;


			// Correct chromatic shift
			N_CHANNELS = nSlices;
			for (j = 0; j < N_CHANNELS; j++) {

				MATRIX_PATH = CHROMA_PATH + "matrix_CH" + CHAN_ARRAY[j] + ".txt";

				stackID = getImageID();
				setSlice(j+1);
				// run("Duplicate...", "duplicate channels=" + (j+1));
				run("Duplicate...", "use");
				dupID = getImageID();
				print("matrix path: " + MATRIX_PATH);
				// in pixels or units?
				run("TransformJ Affine", "matrix=[" + MATRIX_PATH + "] interpolation=[Quintic B-Spline] background=0.0");
				run("Select All");
				run("Copy");
				close();
				selectImage(dupID);
				close();
				selectImage(stackID);
				setSlice(j+1);
				run("Paste");
			}

			// 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 + "_corr-C=" + j;
					rename(NEW_WINDOW_NAME);

					// Create output file path and save the output image
					OUTPUT_PATH = OUTPUT_DIR + NEW_WINDOW_NAME + ".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 + "_corr.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 tif extensions
	}	// end of FOR loop on n extensions


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


	// Restore settings
	restoreSettings();

	setBatchMode("exit and display");

	print("");
	print("*** Correct STORM Channels end ***");
	showStatus("Correct STORM finished");
}


//*************** 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;
}