// Macro Stitch_Mosaic by Christophe Leterrier
// v1.0 05/05/201-6
// v2.0 17/09/2020 added stage position from metadata for Nikon images
// v2.1 18/09/2020 preserves scale for mosaics
// Require S. Preibisch "Grid/Collection Stitching" plugin (included in Fiji)
// Takes an input folder containing single or multi-channel tifs that have been taken as "mosaics" by hand (following an axon for example)
// The input folder can contain several mosaics, and the elements for each must be named by adding "WFa", "WFb", "WFc" / "SrA", "SrB", "SrC" etc. to a common name
// Example: my_cool_mosaic_WFa.tif, my_cool_mosaic_WFb.tif, my_cool_mosaic_WFc.tif, my_other_mosaic_WFa.tif, my_other_mosaic_WFb.tif
// Non mosaic images can exist and be named with "WF" without a,b,c... : my_non_mosaic_image_WF.tif
// v3 23/04/2025 Multichannel 3D with options for global or sequential pairwise stitching and use of Nikon metadata (optional) (version named Stitch Mosaic 3 in WIP)

macro "Stitch_Mosaic" {
	print("\n\n\n***** Stitch Mosaic Log *****");
	print("\n*** Stitch Mosaic has started ***");

// Main parameters
	sort_DEF = true; // Perform images sorting and tiles configuration extraction
	mos_DEF = true;  // Perofrm alignment and stitching	
	batch_DEF = true; // Use batch mode

	MICRO_ARRAY = newArray("NSTORM", "SoRa");
	MICRO_DEF = "SoRa";

	// Use sequential pairwise rather than global, good for a single pair of images (with posdata=false)
	STITCH_ARRAY = newArray("Sequential pairwise", "Global");
	STITCH_DEF = "Sequential pairwise";
		
	posdata_DEF = false; // use stage position in image metadata (use if global stitching of 3+ images)
	
	// Define merging mode (maximum intensity or linear blending, preferred specially for sequential pairwise stitching)
	MERGE_ARRAY = newArray("Max. Intensity", "Linear Blending");
	MERGE_DEF = "Max. Intensity";

	enh_DEF = false; // enhance mosaic after generation (see parameters below)
	closeOut_DEF = true; // close mosaics after generation


// Process in 2D, 3D or 2.5D (3D with no preset Z position, better for 3D pairwise and 3D global stitching)
//	ndim = 2;
	ndim = 2.5;
//	ndim= 3;

// Parameters for the optional image enhancement at the end
	Enhance = false; // Enhance images in addition to fusing them (background subtraction, unsharp mask, contrast enhancement)
	rb_diam = 50; // diameter of the rolling ball for background substraction (50 for a 512x512 image is OK)
	um_mask = 0.3; // weight of the mask for unsharp mark (closer to 1 = sharper)
	cSat = 0.01; // % saturation for contrast enhancement

	setOption("ExpandableArrays", true);

	// Get the folder name
	inDir = getDirectory("Select a directory full of tif images");

	print("input folder:" + inDir);
	
	//	Creation of the dialog box
	Dialog.create("Stitch Mosaic Options");
	Dialog.addCheckbox("Perform image sorting and tiles configuration", sort_DEF);
	Dialog.addCheckbox("Perform mosaic alignment and stitching", mos_DEF);
	Dialog.addCheckbox("Use silent batch mode", batch_DEF);
	Dialog.addChoice("Microscope", MICRO_ARRAY, MICRO_DEF);
	Dialog.addChoice("Stitching mode", STITCH_ARRAY, STITCH_DEF);
	Dialog.addCheckbox("Use Nikon stage coordinates", posdata_DEF);
	Dialog.addChoice("Merging mode", MERGE_ARRAY, MERGE_DEF);
	Dialog.addCheckbox("Enhance image after mosaic generation", enh_DEF);
	Dialog.addCheckbox("Close mosaic after generation", closeOut_DEF);
	Dialog.show();

//	Feeding variables from dialog choices
	sort = Dialog.getCheckbox();
	mos = Dialog.getCheckbox();
	batch = Dialog.getCheckbox();
	MICRO = Dialog.getChoice();
	STITCH = Dialog.getChoice();
	posdata = Dialog.getCheckbox();
	mergeMode = Dialog.getChoice();
	Enhance = Dialog.getCheckbox();
	closeOut = Dialog.getCheckbox();

	if (batch == true)	setBatchMode(true);
	
	// Get name of input folder, parent folder, short name of input folder (before first space in name)
	parDir = File.getParent(inDir);
	inName = File.getName(inDir);
	if (sort == false) inName = replace(inName, " folders", ""); // to start with sorted folders name when you only want to do the alignment and fusion step
	inShortA = split(inName, " ");
	inShort = inShortA[0];

	// Microscope parameters: pixel size, voxel height, camera full sensor size...
	if (MICRO == "NSTORM") {
		TileID = "WF";
		pixSize = 0.16;
		ZSize = 0.16;
		fullWidth = 512;
		Xfactor = 1; Yfactor = -1; Zfactor = 1;
	}
	else {
		TileID = "Sr";
		pixSize = 0.027;
		ZSize = 0.15; 
		fullWidth = 2304;
		Xfactor = -1; Yfactor = 1; Zfactor = 1;
	}	

	if (STITCH == "Sequential pairwise") pairwise = true;
	else pairwise = false;

	
	if (sort == true) {	
		print("\n*** Sorting and tiles configuration step ***");
		
		// Count the number of files ending with ".tif" as the number of images
		allNames = getFileList(inDir);
		Array.sort(allNames);
		imN = 0;
		for (i = 0; i < allNames.length; i++) {
			nLength = lengthOf(allNames[i]);
			if (substring(allNames[i], nLength-4, nLength) == ".tif") {
				imN ++;
			}
		}
		print("Number of .tif images in input folder:" + imN);
	
		// Store the images names in an array
		imNames = newArray(imN);
		for (i = 0; i < allNames.length; i++) {
			nLength = lengthOf(allNames[i]);
			if (substring(allNames[i], nLength-4, nLength) == ".tif") {
				imNames[i] = allNames[i];
			}
		}
	
	// Create a new folder that will have each mosaic element (individual multi-channel stacks) inside its own subfolder (necessary for the fusion plugin)
	// Mosaic elements name must contain WFa, WFb, WFc etc. preceded by the SAME string.
	// If some images are single (not part of a mosaic) they must contain "WF_" or "WF." in their name.
	
		// Make temp folder where individual folders will contain mosaic elements (single multi-channel stacks)
		tempDir = parDir + File.separator + inName + " folders" + File.separator;
		print("Sorted folder: " + tempDir);
		if (File.isDirectory(tempDir) == false) {
			File.makeDirectory(tempDir);
		}
	
		mosCount = 0; // counter of mosaic number

		// Loop on images names
		for (i = 0; i < imN; i++) {
			// If image is first in a mosaic series (name has WFa, followers have WFb, WFc etc...)
			WFi = indexOf(imNames[i], TileID + "a");
			WFi2 = indexOf(imNames[i], TileID + "A");
			// If image is isolated (not part of mosaic) it has "WF_" or "WF." in its name
			WFn = indexOf(imNames[i], TileID + ".");
			WFc = indexOf(imNames[i], TileID + "_");
			// Store mosaic elements / individual images in separated folders
			if (WFi > -1 || WFi2 > -1 || WFn > -1 || WFc > -1) {
			// Create a subfolder of temp folder
				fName = substring(imNames[i], 0, WFi + WFi2 + WFn + WFc + 5);
				// case of a mosaic: the folder name will end with "WFm"
				if (WFi > -1 || WFi2 > -1) subName = fName + "m";
				// case of single image: the folder name will end with "WF"
				else subName = fName;
				// make the subfolder
				fDir = tempDir + subName + File.separator;
				if (File.isDirectory(fDir) == false) {
					File.makeDirectory(fDir);
				}
			// Loop on all images in input folder to find all parts of the mosaic (including the "WFa" or "WF_" / "WF." one)
				
				// save the mosaic elements 
				flag2 = 0; // flag
				
				for (j = 0; j < imN; j++) {
					if (indexOf(imNames[j], fName) > -1) { // only for images of this mosaic
									
						// open image
						open(inDir + imNames[j]);
						
						// if first image from mosaic, get scale and store in array
						flag2++; // flag to get scale at first image of the mosaic
						if (flag2 == 1) {
							mosCount++;	// mosaic counter
						}
						
						save(fDir + imNames[j]);
						close();
					}
				}
			
			print("Mosaic #" + mosCount + " (" + flag2 + " tiles" + "): " + fDir);
			
			}
		}
	
		// print("Total number of mosaics: " + mosCount);
	
		// Generate TileConfiguration.txt file from image metadata if option checked
		
		if (posdata == true) {
		//Loop on temp subfolders to process mosaics
		foldNames = getFileList(tempDir);
		for (i = 0; i < foldNames.length; i++) {
			if (indexOf(foldNames[i], TileID + "m") > -1) { // only if not single image
					// call the function to make the configuration file from stage position metadata in images (Nikon version)
					makeTileConfiguration(tempDir + foldNames[i]);
				}	
			}
		}
	}

	if (mos == true) {

		// Perform stiching using S. Preibisch's "Grid/Collection Stitching" plugin
		// part of Fiji see http://imagej.net/Image_Stitching#Grid.2FCollection_Stitching
		
		print("\n*** Alignment and stitching step ***");
		
		// Useful if macro just makes alignment form existing sorted folder
		if (sort == false) {
			tempDir = parDir + File.separator + inName + " folders" + File.separator;
			print ("Sorted tiles folder: " + tempDir);
		}
		
		// Create mosaics folder
		mosDir = parDir + File.separator + inName + " mosaics" + File.separator;
		print("Mosaic folder: " + mosDir);
		if (File.isDirectory(mosDir) == false) {
			File.makeDirectory(mosDir);
		}

		//Loop on temp subfolders to process mosaics
		foldNames = getFileList(tempDir);
		for (i = 0; i < foldNames.length; i++) {
			// Stitch all images in subfolder
			if (indexOf(foldNames[i], TileID + "m") > -1) {
				
				// open first image to get spatial scale
				tileNames = getFileList(tempDir + foldNames[i]);
				tileNames = Array.sort(tileNames);
				open (tempDir + foldNames[i] + tileNames[0]);
				getPixelSize(scaleUnit, pixelWidth, pixelHeight);
				close();
								
				// if using stage positions in images metadata 
				if (posdata == true) {
					// call the function to make the configuration file from stage position metadata in images (Nikon version)
					if (pairwise == false) {
						print("--- Running global stitching with preset positions on " + foldNames[i]);
						run("Grid/Collection stitching", "type=[Positions from file] order=[Defined by TileConfiguration] directory=[" + tempDir + foldNames[i] +"] layout_file=TileConfiguration.txt fusion_method=[" + mergeMode + "] regression_threshold=0.30 max/avg_displacement_threshold=2.50 absolute_displacement_threshold=3.50 compute_overlap subpixel_accuracy computation_parameters=[Save computation time (but use more RAM)] image_output=[Fuse and display]");
					}
					else {
						print("--- Running sequential pairwise stitching with preset positions on " + foldNames[i]);
						seqPair(tempDir + foldNames[i], 1);
					}
				}
				
				else {
					if (pairwise == false) {
						print("--- Running global stitching without preset positions on " + foldNames[i]);
						run("Grid/Collection stitching", "type=[Unknown position] order=[All files in directory] directory=[" + tempDir + foldNames[i] +"] output_textfile_name=TileConfiguration.txt fusion_method=[" + mergeMode + "] regression_threshold=0.30 max/avg_displacement_threshold=2.50 absolute_displacement_threshold=3.50 ignore_z_stage subpixel_accuracy computation_parameters=[Save computation time (but use more RAM)] image_output=[Fuse and display]");
					}
					else {
						print("--- Running sequential pairwise stitching without preset positions on " + foldNames[i]);
						seqPair(tempDir + foldNames[i], 0);
					}
				}
			}
			// unless it's a single image
			else {
				print("--- No stitching for single image " + foldNames[i]);
				insideNames = getFileList(tempDir + foldNames[i]);
				open(insideNames[0]);
				if (nSlices > 1) Stack.setDisplayMode("composite");
			}
			// Enhance fused image (contrast enhancement looping on channels)
			if (Enhance == true) {
				print("     Enhancing fused image");
				getDimensions(w, h, ch, sl, fr);
				for (j = 0; j < ch; j++) {
					if (nSlices > 1) Stack.setChannel(j+1);
					run("Enhance Contrast...", "saturated=" + cSat);
				}
			}
	
			// Apply stored scale
			run("Set Scale...", "distance=1 known=" + pixelWidth + " pixel=1 unit=" + scaleUnit);
			
			// Save resulting image
			fusedName = substring(foldNames[i], 0, lengthOf(foldNames[i]) - 1) + "_fused.tif";
			rename(fusedName);
			print("    Saved fused image as " + mosDir + fusedName);
			save(mosDir + fusedName);
			if (closeOut == true) close();
		}
	}
	
	if (batch == true) setBatchMode("exit and display");
	print("\n*** Stitch Mosaic has finished ***");
}

function makeTileConfiguration(InputDir) {
	
	//	Get all file names
	AllNames = getFileList(InputDir);
	Array.sort(AllNames);
	nL = AllNames.length;
	AllExt = newArray(nL);
	
	//	Create extensions array and test if tif
	ntif = 0;
	for (k = 0; k < nL; k++) {
		AllNamesParts = getFileExtension(AllNames[k]);
		AllExt[k] = AllNamesParts[1];
		if (indexOf(toLowerCase(AllExt[k]), ".tif") > -1) ntif++; // condition to catch .tif, .tiff, .TIF, .TIFF...
	}

	// Define arrays for metadata of each tile
	NamesArray = newArray(ntif);
	XPosArray = newArray(ntif);
	YPosArray = newArray(ntif);
	leftArray = newArray(ntif);
	topArray = newArray(ntif);
	sizeXArray = newArray(ntif); 
	sizeYArray = newArray(ntif); 
	if (ndim == 3) ZPosArray = newArray(ntif);

	// Get the coordinates for each tile from the metadata
	for (l = 0; l<nL; l++) {
		if (indexOf(toLowerCase(AllExt[l]), ".tif") > -1) { // open only tif files
			open(InputDir + AllNames[l]);
			
			// get X and Y stage position in µm, top and left coordinates of the sensor crop in pixels, size of image in pixels, Z coordinate in µm (Nikon metadata)
			if (ndim != 3) XYCoor = getPropValues("dXpos,dYPos,left,top,uiWidth,uiHeight");
			else XYCoor = getPropValues("dXpos,dYPos,left,top,uiWidth,uiHeight,dZPos");

			// Assign name of the image, X and Y stage position, left/top coordinates, image size, Z coordinate in arrays
			NamesArray[l] = AllNames[l];
			XPosArray[l] = XYCoor[0];
			YPosArray[l] = XYCoor[1];
			leftArray[l] = XYCoor[2];
			topArray[l] = XYCoor[3];
			sizeXArray[l] = XYCoor[4];
			sizeYArray[l] = XYCoor[5];
			if (ndim == 3) ZPosArray[l] = XYCoor[6];
			
//			print("NameArray[" + l + "]: " + NamesArray[l]); 
//			print("   XPosArray[" + l + "]=" + XPosArray[l] + ", YPosArray[" + l + "]=" + YPosArray[l]);
//			print ("   leftArray[" + l + "]=" + leftArray[l] + ", topArray[" + l + "]=" + topArray[l] + ", sizeXArray[" + l + "]=" + sizeXArray[l] + ", sizeYArray[" + l + "]=" + sizeYArray[l]); 
//			if (ndim == 3) print ("   ZPosArray[" + l + "]=" + ZPosArray[l]);

			// close image
			close();
		}
	}
	
	// Define intermediate and output array
	Xcoor = newArray(ntif);
	Ycoor = newArray(ntif);
	XArrayT = newArray(ntif);
	YArrayT = newArray(ntif);
	Zcoor = newArray(ntif);
	ZArrayT = newArray(ntif);
		
	// Calculate and store the center of each tile
	for (m = 0; m<ntif; m++) {		
		// Center of the sensor for each tile in pixels
		Xcenter = Xfactor * (parseFloat(XPosArray[m]) - parseFloat(XPosArray[0])) / pixSize;
		Ycenter = Yfactor * (parseFloat(YPosArray[m]) - parseFloat(YPosArray[0])) / pixSize;
	 
		
		// Shift to account for non-centered crop
		Xcoor[m] = Xcenter + parseFloat(leftArray[m]) - (fullWidth/2) + (parseFloat(sizeXArray[m])/2);
		Ycoor[m] = Ycenter + parseFloat(topArray[m]) - (fullWidth/2) + (parseFloat(sizeYArray[m])/2);
		
		// Calculate Z position in pixels
		if (ndim == 3) Zcoor[m] = Zfactor * (parseFloat(ZPosArray[m]) - parseFloat(ZPosArray[0])) / ZSize;
		else Zcoor[m] = 0;
		
		// Rounding to 1 decimal
		XArrayT[m] = toString(Xcoor[m], 1);
		YArrayT[m] = toString(Ycoor[m], 1);
		if (ndim == 3) ZArrayT[m] = toString(Zcoor[m], 1);
		else ZArrayT[m] = 0;
		
		print("   Image #" + (m+1) + ": " + NamesArray[m]);
		print ("     width: " + sizeXArray[m] + ", heigth: " + sizeYArray[m]);
		print("     X position: " + XArrayT[m] + ", Y position: " + YArrayT[m]);
		if (ndim == 3) print("     Z position: " + ZArrayT[m]);
		
	}
	
	// Write TileConfiguration.txt file
	writeTileConfiguration(NamesArray, XArrayT, YArrayT, ZArrayT, InputDir);
	
	return;

}

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


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 (o = 0; o < pL; o++) {
		
		// split each line according to delimiters defined above (del1, del2) or assign NaN if no delimiter detected
		if (indexOf(InfoArray[o], del1) > -1) LineSplit = split(InfoArray[o], "(" + del1 + ")"); // the split string is parenthesed to work as regular expression
			else if (indexOf(InfoArray[o], del2) > -1) LineSplit = split(InfoArray[o], "(" + del2 + ")");
			else LineSplit = newArray(NaN, NaN);
	
		PropArray[o] = LineSplit[0];
		ValArray[o] = 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 (p = 0; p < kL; p++) {
		flag = 0;
		for (q = 0; q < pL; q++) {
			if (PropArray[q] == KeyArray [p]) {
				flag = 1;
				KeyValArray[p] = ValArray[q];
			}
		}
		if (flag == 0) KeyValArray[p] = NaN; // if key not found, assign NaN for that slot
	}

	// return the key values array
	return KeyValArray;
}


function writeTileConfiguration(nameA, xA, yA, zA, fp) {
	
	// beginning of the file
	if (ndim == 2) FileString = "# Define the number of dimensions we are working on\ndim = 2\n\n# Define the image coordinates\n";
	else FileString = "# Define the number of dimensions we are working on\ndim = 3\n\n# Define the image coordinates\n";
	
	// lines for each image and its coordinates
	for (r = 0; r < nameA.length; r++) {
		if (ndim == 2) FileString += nameA[r] + "; ; (" + xA[r] + ", " + yA[r] + ")\n";	
		else if (ndim == 2.5) FileString += nameA[r] + "; ; (" + xA[r] + ", " + yA[r] + ", 0.0)\n";
		else FileString += nameA[r] + "; ; (" + xA[r] + ", " + yA[r] + ", " + zA[r] + ")\n";
	}

	// save file
	File.saveString(FileString, fp + "TileConfiguration.txt");
	
	return;
	
}

function seqPair(ip, po) { // takes the folder path and a position option (1=preset positions, 0=no preset positions)
	
	// Make an array of tif images names in folder
	pairFileNames = getFileList(ip);
	pairFiles = newArray;
	t = 0;
	for (s=0; s<pairFileNames.length; s++) {
		fileNameParts = getFileExtension(pairFileNames[s]);
		fileExt = fileNameParts[1];
		if (indexOf(toLowerCase(fileExt), ".tif") > -1) {
			pairFiles[t] = pairFileNames[s];
			t++;
		}
	}
	
	if (po == 1) {
		presCoords = getPresetPos(ip);	
		presX = newArray(presCoords.length);
		presY = newArray(presCoords.length);
		if (ndim != 2) presZ = newArray(presCoords.length);
		
		for (w=0; w<presCoords.length; w++) {
			splitPresCoords = split(presCoords[w], "(, )");
			presX[w] = splitPresCoords[0];
			presY[w] = splitPresCoords[1];
			if (ndim != 2) presZ[w] = splitPresCoords[2];
			else presZ[w] = 0;
		}
	}
	
	open(ip + pairFiles[0]);
	firstID = getImageID();
	firstTitle = getTitle();
	
	
	for (u=1; u<pairFiles.length; u++) {
		
		open(ip + pairFiles[u]);
		secondID = getImageID();
		secondTitle = getTitle();
		
		if (po == 0) posString = "x=0.0000 y=0.0000 z=0.0000";
		else posString = "x="+ presX[u] + " y=" + presY[u] + " z=" + presZ[u];
		
		print("fusing with " + secondTitle);
		if (posdata == true) print("using preset positions: " + posString);
		else print ("without preset positions");
		run("Pairwise stitching", "first_image=" + firstTitle + " second_image=" + secondTitle +  " fusion_method=[" + mergeMode + "] fused_image=fused.tif check_peaks=5 compute_overlap subpixel_accuracy " + posString + " registration_channel_image_1=[Average all channels] registration_channel_image_2=[Average all channels]");
		fusedID = getImageID();
		fusedTitle = getTitle();		
		selectImage(firstID);
		close();
		selectImage(secondID);
		close();
		firstID = fusedID;
		firstTitle = fusedTitle;
	}
	
	return;	
}


function getPresetPos(inp) {
	
	presetString = File.openAsString(inp + "TileConfiguration.txt");
	presetLines = split(presetString, "\n");
	pairCoords = newArray(presetLines.length-4);
	
	for (v=4; v<presetLines.length; v++) {
		openPar = lastIndexOf(presetLines[v], "(");
		closePar = lastIndexOf(presetLines[v], ")");
		pairCoords[v-4] = substring(presetLines[v], openPar+1, closePar);
	}
	
	return pairCoords;	
}