// Calibrate_Chroma macro by Christophe Leterrier
// 11/04/2022
//
// Processes a folder full of multichannel tif, register each channel relative to one of them using the NanoJ-core plugin;
// then generates an average translation mask from each calculated translation mask
//
// Correction of stacks is then made using the Correct_Apo_2 macro, applying the NanoJ-Core translation mask correction

macro "Calibrate_Chroma" {


// Default values for the Options Panel
	CROP_DEF = "Full"; // Center crop option
	FILT_DEF = true; // Fourier filtering
	REF_CH_DEF = 1; // means reference channel is first channel ie C=0
	SIM_DEF = 0.5; // Minimal similarity (from tests: 0.5 for apo, 0.25 for STORM 256, 0.1 for STORM 512)
	GAUSS_DEF = 1; // gaussian blur before aligning

// Initialize choices variables
	CROP_ARRAY = newArray("Full", "1024", "512", "256");
	
// Input folder
	INPUT_DIR=getDirectory("Select the input stacks directory");
	ALL_NAMES=getFileList(INPUT_DIR);
	Array.sort(ALL_NAMES);

	print("\n\n\n*** Calibrate Chroma Log ***");
	print("");

// Options
	Dialog.create("Calibrate Chroma Options");
	
	Dialog.addChoice("Crop to center pixels:", CROP_ARRAY, CROP_DEF);
	Dialog.addCheckbox("Filter background", FILT_DEF);
	Dialog.addNumber("Reference channel:", REF_CH_DEF);
	Dialog.addNumber("Minimal similarity:", SIM_DEF, 2, 4, "");
	Dialog.addNumber("Gaussian blur:", GAUSS_DEF, 1, 3, "px");
	Dialog.show();

	// Feeding variables from dialog choices
	CROP = Dialog.getChoice();
	FILT = Dialog.getCheckbox();
	REF_CHANNEL = Dialog.getNumber();
	SIM = Dialog.getNumber();
	GAUSS = Dialog.getNumber();

	setBatchMode(true);

	print("Processing images in folder:" + INPUT_DIR);

	// Stack counter for average
	f = 0;
	
	// Array of trasnlation mask titles
	TM_TITLES = newArray(0);
	CONCAT_STRING = "open "; 

	// Loop on tif images
	for (n=0; n<ALL_NAMES.length; n++) {
		NAME_PARTS = split(ALL_NAMES[n], ".");
		NAME_EXT = NAME_PARTS[NAME_PARTS.length - 1];

		if (NAME_EXT == "tif") {

			// increment stack counter
			f++;

			// Open stack
			open(INPUT_DIR + ALL_NAMES[n]);
			IN_TITLE = getTitle();
			IN_ID = getImageID();
			
			print("    Processing image #" + f + ": " + IN_TITLE);
			
			// Create flags & sum arrays at first stack
			getPixelSize(unit, pw, ph, pd);
			N_CHANNELS = nSlices;
			SIZEW = getWidth();
			SIZEH = getHeight();

			print("    Image has " + N_CHANNELS + " channels");			
			print("    Image is " + SIZEW + "x" + SIZEH + " pixels");
			print("    Pixel size is " + pw + " " + unit);
			

			
			// Crop to center
			if (CROP != "full") {
				CENTERX = Math.floor(SIZEW/2);
				CENTERY = Math.floor(SIZEH/2);
			
				run("Specify...", "width=" + CROP + " height=" + CROP + " x=" + CENTERX + " y=" + CENTERY + " centered");	
				run("Crop");
				IN_TITLE = getTitle();
				IN_ID = getImageID();
				
				print("        ...cropped");	
			}

			// Filter low frequencies (background)
			if (FILT == true) {
				run("Duplicate...", "duplicate");
				DUP_TITLE = getTitle();
				DUP_ID = getImageID();
				FLARGE = SIZEW;
				FSMALL = Math.floor(SIZEW/4);
				run("Bandpass Filter...", "filter_large=" + FLARGE + " filter_small=" + FSMALL + " suppress=None tolerance=5 process");
				imageCalculator("Subtract stack", IN_TITLE, DUP_TITLE);
				IN_TITLE = getTitle();
				IN_ID = getImageID();
				selectImage(DUP_ID);
				close();
				print("        ...filtered");
			}
			
			// Convert into regular stack to use NanoJ-Core
			run("Hyperstack to Stack");
			IN_TITLE = getTitle();
			IN_ID = getImageID();

			// Run NanoJ-Core chromatic aberration correction estimation
			run("Register Channels - Estimate", "reference=" + REF_CHANNEL + " number=" + N_CHANNELS + " max=0 blocks=5 min=" + SIM + " gaussian=" + GAUSS);
			// Convert output to multi-channel hyperstack
			run("Stack to Hyperstack...", "order=xyczt(default) channels=" + N_CHANNELS + " slices=1 frames=1 display=Grayscale");	
			CURR_TITLE = getTitle();
			// Add resulting mask title to the list of titles
			Array.concat(TM_TITLES, CURR_TITLE);
			CONCAT_STRING += " image" + f + "=[" + CURR_TITLE + "]";
			
			// close input image
			selectImage(IN_ID);
			close();
			
			print("        ...aligned");

		} // end of if loop on tif images
	} // end of for loop on all files

	
	run("Concatenate...", CONCAT_STRING);
	CONC_TITLE = getTitle();
	CONC_ID = getImageID();
	run("Enhance Contrast", "saturated=0.35");
	run("Z Project...", "projection=[Average Intensity]");
	
	OUT_PATH = INPUT_DIR + "TranslationMask.tif";
	save(OUT_PATH);
	print("Translation mask (averaged over " + f + " images) saved as " + OUT_PATH);
	close();
	close();
	
	setBatchMode("exit and display");

	print("\n*** Calibrate_Chroma end ***");
	showStatus("Calibrate_Chroma finished");
}




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

function getLogName(Name) {
	tifindex = indexOf(Name, ".tif");
	zviindex = indexOf(Name, ".zvi");
	if (zviindex > 0) index = minOf(tifindex, zviindex);
	else index = tifindex;
	short = substring(Name, 0, index);
//	print(short);
	return short;
}