// Batch_3Ddrift_ZOLA macro by Christophe Leterrier
// Process ThunderSTORM csv localization files
// Drift correction in ZOLA, export, then reconversion into TS format

macro "Batch 3D Drift ZOLA" {

	// Save Settings
	saveSettings();

//*************** Initialize variables ***************	

	
	// Detect if called from macro
	arg = getArgument();
	if (lengthOf(arg)>0) {
		called = true;
		argarray = split(arg, ",");
	}
	else {
		called = false;
	}

	LOC_SUFFIX1 = ".csv";
	LOC_SUFFIX2 = ".xls";
	OUT_PARAM = "procZ";
	
	// Camera setup variables
	CAM_SIZE = 160;
	CG = 12.48; // camera gain new images
	EM = 100; // EM gain new images	

	// File chooser
	CHOOSE_DEF = false;
	CHOOSE_STRING_DEF = "";

	EXC_DEF = true;
	EXC_STRING_DEF = "_ZR_";

	COUNT_DEF = true;

	// Drift correction
	CORR_DRIFT_DEF = true;
	BIN_DEF = 7; // number of frames per sub-reconstruction used for autocorrelation
	MAG_DEF = 50; //  pixel size (nm) of sub-reconstructions used for autocorrelation
	SM_DEF = 1; // Maximal drift (um)
	
//*************** Get input folder ***************	

	// Get input directory (dialog or argument)
	if (called == false) {
		INPUT_DIR = getDirectory("Select a source folder");
	}
	else {
		INPUT_DIR = argarray[0];
	}
	
	print("\n\n\n*** Batch 3D Drift ZOLA started ***");
	print("");
	print("Input folder: " + INPUT_DIR);


//*************** Dialog ***************	

	if (called == false) {	
		//Creation of the dialog box
		Dialog.create("Batch process localizations: options");
		Dialog.addCheckbox("Choose files based on name", CHOOSE_DEF);
		Dialog.addString("Name contains", CHOOSE_STRING_DEF);
		Dialog.addCheckbox("Exclude files based on name", EXC_DEF);
		Dialog.addString("Name contains", EXC_STRING_DEF);
		Dialog.addCheckbox("Reset Loc Counter", COUNT_DEF);
		Dialog.addMessage("");
		Dialog.addCheckbox("Correct drift", CORR_DRIFT_DEF);
		Dialog.addNumber("Number of bins for sub-images", BIN_DEF, 0, 5, "");
		Dialog.addNumber("Magnification for autocorrelation", MAG_DEF, 0, 2, "X");
		Dialog.addNumber("Smoothing factor for autocorrelation", SM_DEF, 3, 5, "X");
		Dialog.show();
		
		// Feeding variables from dialog choices
		CHOOSE = Dialog.getCheckbox();
		CHOOSE_STRING = Dialog.getString();

		EXC = Dialog.getCheckbox();
		EXC_STRING = Dialog.getString();

		RESET = Dialog.getCheckbox();
		
		CORR_DRIFT = Dialog.getCheckbox();
		BIN = Dialog.getNumber();
		MAG = Dialog.getNumber();
		SM = Dialog.getNumber();
	}

	// called from macro:
	// arguments (INPUT_DIR, CHOOSE, CHOOSE_STRING, EXC, EXC_STRING, CORR_DRIFT, BIN, MAG, SM)
	else {
		CHOOSE = argarray[1];
		CHOOSE_STRING = argarray[2];
		EXC = argarray[3];
		EXC_STRING = argarray[4];
		RESET = argarray[5];		
		CORR_DRIFT = argarray[6];
		BIN = argarray[7];
		MAG = argarray[8];
		SM = argarray[9];
	}
	
//*************** Prepare processing ***************	

	//Time counter
	startTime = getTime();
	
	// Get all file names
	ALL_NAMES = getFileList(INPUT_DIR);
	Array.sort(ALL_NAMES);
	ALL_TS = newArray(ALL_NAMES.length);
	ALL_TYPES = newArray(ALL_NAMES.length);
	
	OUTPUT_DIR = File.getParent(INPUT_DIR);
	OUTPUT_NAME = File.getName(INPUT_DIR);
	OUTPUT_DIR = OUTPUT_DIR + File.separator + OUTPUT_NAME + " " + OUT_PARAM + File.separator;

	if (File.isDirectory(OUTPUT_DIR) == false) {
		File.makeDirectory(OUTPUT_DIR);
	}

	print("Output folder: " + OUTPUT_DIR);

	run("EMCCD", "emccd_camera_adu=" + CG + " emccd_camera_gain=" + EM + " emccd_camera_offset=100");

//*************** Process loc files ***************

	// Detect number of files
	FileTotal = 0;
	for (n = 0; n < ALL_NAMES.length; n++) {
		ALL_TS[n] = false;
		ALL_TYPES[n] = "not TS";
		if (endsWith(ALL_NAMES[n], LOC_SUFFIX1) == true) {
			LOC_SUFFIX = LOC_SUFFIX1;
			ALL_TS[n] = true;
			ALL_TYPES[n] = "[CSV (comma separated)]";
		}
		if (endsWith(ALL_NAMES[n], LOC_SUFFIX2) == true) {
			LOC_SUFFIX = LOC_SUFFIX2;
			ALL_TS[n] = true;
			ALL_TYPES [n] = "[XLS (tab separated)]";
		}		
		if (ALL_TS[n] == true) {
			if ((CHOOSE == false || indexOf(ALL_NAMES[n], CHOOSE_STRING) > -1) && (EXC == false || indexOf(ALL_NAMES[n], EXC_STRING) == -1)) {
				FileTotal++;
			}
		}
	}

	// print(FileTotal);

	// Loop on all TS loc files
	FileCount = 0;
	for (n = 0; n < ALL_NAMES.length; n++) {
		if (ALL_TS[n] == true) {
			if ((CHOOSE == false || indexOf(ALL_NAMES[n], CHOOSE_STRING) > -1) && (EXC == false || indexOf(ALL_NAMES[n], EXC_STRING) == -1)) {
				// Image counter
				FileCount++;
				
				// Get the file path
				FILE_PATH = INPUT_DIR + ALL_NAMES[n];
				
				// Store components of the file name
				FILE_NAME = File.getName(FILE_PATH);
				
				print("    Input file #" + FileCount + "/" + FileTotal + ": " + FILE_NAME);
	
				OUT_TITLE = FILE_NAME;
				
				// Open the loc file	
				run("Import table", "file_path=[" + FILE_PATH +]");
	
				
				if (CORR_DRIFT == true) {
					run("3D Drift correction", "run_on_gpu cross-correlation_pixel_size=" + MAG + " number=" + BIN + " maximum_drift=" + SM + " localization_table_attached=[]");
				}

				// Export the corrected locs into an output file












				// Count the new Loc number
				nLocs = eval("script", "importClass(Packages.cz.cuni.lf1.lge.ThunderSTORM.results.IJResultsTable); var rt = IJResultsTable.getResultsTable(); rows = rt.getRowCount();");
				nLocK = round(nLocs / 1000);
				// OUT_TITLE = replace(OUT_TITLE, "([0-9])+K_", nLocK + "K_");

				//if (CORR_DRIFT == true) {
				//	ADD_TITLE = ADD_TITLE + "DC";
				//}		

				ADD_TITLE = "_" + nLocK + "K";
				
				if (indexOf(OUT_TITLE, "_TS2D.") > 0 || indexOf(OUT_TITLE, "_TS3D.") > 0) {
					
					if (RESET == true) {
						NEW_TITLE = replace(OUT_TITLE, "(_([0-9])+K)+_TS", ADD_TITLE + "_TS");
					}
					else {
						NEW_TITLE = replace(OUT_TITLE, "_TS", ADD_TITLE + "_TS");
					}
				}

				else {
					NEW_TITLE = replace(OUT_TITLE, LOC_SUFFIX, ADD_TITLE + LOC_SUFFIX1);
				}

				
/*
 				else {
					if (RESET == true) {
						NEW_TITLE = replace(OUT_TITLE, "(_([0-9])+K)+", ADD_TITLE);
					}
					else {
						insert = lastIndexOf(OUT_TITLE, "K_");
						NEW_TIT1 = substring(OUT_TITLE, 0, insert+1);
						NEW_TIT2 = substring(OUT_TITLE, insert+1, lengthOf(OUT_TITLE));
						NEW_TITLE = NEW_TIT1 + ADD_TITLE + NEW_TIT2;
					}				
				}
*/
				

				OUT_PATH = OUTPUT_DIR + NEW_TITLE;
				run("Export results", "filepath=[" + OUT_PATH + "] fileformat=[CSV (comma separated)] chi2=false saveprotocol=false");
				// remove chi2 that causes an error on 27-07-2017 version (see bug on GitHub)
				
				
				print("      Output file:" + OUT_TITLE);
	
				// Rename the drift image
				if (CORR_DRIFT == true) {
					selectWindow("Drift");
					rename(OUT_TITLE);
				}
				
			} // end of IF loop on include/exclude names	
		}	// end of IF loop on extensions
	}	// end of FOR loop on n extensions

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


	if (CORR_DRIFT == true && FileCount > 1) {
			run("Images to Stack", "name=[Drift Correction Stack] title=" + ".csv" + " use");
			save(OUTPUT_DIR + "Drift.tif");
			close();
	}

	//NEW_INPUT = substring(INPUT_DIR, 0, lengthOf(INPUT_DIR) - 1) + " no proc" + File.separator;
	//File.rename(INPUT_DIR, NEW_INPUT);
	//File.rename(OUTPUT_DIR, INPUT_DIR);
	
	// Restore settings
	restoreSettings();	
	showStatus("Batch Process Localizations finished");
	
	//Time counter
	stopTime = getTime();
	Time = stopTime - startTime;
	
	print("");
	print("*** Batch 3D Drift ZOLA end after " + Time / 1000 + " s ***\n\n\n");
	return OUTPUT_DIR;
}