// F-TSAddZUnc.js script function by Christophe Leterrier
// Add (or replace) Z uncertainty to a ThunderSTORM file by multiplying the ZY uncertainty by a factor f (usually 2);
// 14-10-26

importClass(Packages.java.io.File);
importClass(Packages.java.io.FileReader);
importClass(Packages.java.io.FileWriter);
importClass(Packages.java.io.BufferedReader);
importClass(Packages.java.io.BufferedWriter);
importClass(Packages.ij.IJ);
importClass(Packages.java.lang.Double);

function TSAddZUnc(inPath, outDir, f) {

	// separators (csv files)
	var inSep = ",";
	var sep = ",";

	var zHeader = "\"z [nm]\"";
	var xyUncHeader = "\"uncertainty_xy [nm]\"";
	var zUncHeader = "\"uncertainty_z [nm]\"";

	var SampleMax = 100;

	// Define input files, folder, open it etc.
	var inFile = new File(inPath); 
	var inName = inFile.getName();
	var inNameExt = getExt("" + inName);
	IJ.log("    inName: " + inName);
	
	var br = new BufferedReader(new FileReader(inFile));

	// Get the header line and find index of Z column
	inHLine = br.readLine();
	IJ.log("      inHLine: " + inHLine);	
	var inHArray = inHLine.split(inSep);
	var zIndex = arrayFind(inHArray, zHeader);
	
	var zAccu = 0;
	var i = 0;	
	while ((inLine = br.readLine()) != null && i < SampleMax) {
		i++;
		inCells = inLine.split(inSep);
		zAccu = zAccu + Double.parseDouble(inCells[zIndex]);
	}

	// Close the input file and reopens it (back to line 1), reads header line
	br.close();
	var br = new BufferedReader(new FileReader(inFile));
	var inLine = br.readLine();

	// Generate output name and path, open file writer
	var outName = inName.replace("TS3D", "unc_TS3D");
	if (outName == inName) outName = inNameExt[0] + "_unc." + inNameExt[1];
	var outPath = outDir + outName;
	var outFile = new File(outPath);

	if (!outFile.exists()) {
		outFile.createNewFile();
	}
	var bw = new BufferedWriter(new FileWriter(outFile));

	IJ.log("      outName: " + outName);

	// Detect XY and Z uncertainty columns
	var xyUncIndex = arrayFind(inHArray, xyUncHeader);
	var zUncIndex = arrayFind(inHArray, zUncHeader);

	// Warn about 2D file
	if (zAccu == 0) {
		IJ.log("      2D file: Z uncertainty set to 0");
	}
	// Add Z uncertainty column only if not already existing
	if (zUncIndex == -1) {
		IJ.log("      adding Z uncertainty");
		inHLine = inHLine + sep + "uncertainty_z [nm]";
		bw.write(inHLine);
	}
	else {
		IJ.log("      replacing Z uncertainty");
		bw.write(inHLine);
	}
	
	bw.newLine();

	// Write the output file line by line
	while ((inLine = br.readLine()) != null) {

		inCells = inLine.split(inSep);
		
		var xyUnc = Double.parseDouble(inCells[xyUncIndex]);
		if (zAccu == 0) var zUnc = 0;
		else var zUnc = f * xyUnc;
		
		// Append Z uncertainty value if not existing
		if (zUncIndex == -1) {
			var outLine = inCells[0];
			for (var t= 1; t < inCells.length; t++) {
				outLine = outLine + sep + inCells[t]; 
			}
			outLine = outLine + sep + zUnc; 
		}

		// or replace the existing value
		else {
			inCells[zUncIndex] = zUnc.toString();			
			var outLine = inCells[0];
			for (var t= 1; t < inCells.length; t++) {
				outLine = outLine + sep + inCells[t]; 
			}
				
		}
		bw.write(outLine);
		bw.newLine();
	}		 
	br.close();
	bw.close();
}


function getExt(filestring){
	var namearray = filestring.split(".");
	var shortname = "";
	for (var f = 0; f < namearray.length - 1; f++) {
		shortname = shortname + namearray[f];
	}
	return [shortname, namearray[namearray.length - 1]];
}

function arrayFind(a, s){
	index = -1;
	for (var i = 0; i < a.length; i++) {
		if (a[i] == s) index = i;	
	}
	return index;
}