// F-TSWarp.js script function by Christophe Leterrier
// Warps a source channel localization file according to the warps coeeficients as found in a Nikon N-STORM.ini file (see Batch_Warp_Channels.js)
// 15-07-16

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);


/*
Warps are always longer WL to shorter WL ie 647 -› 488 for example
"Source WL" is the WL corresponding to the image that will be transformed (647)
"Destination WL" is the target WL to which the source WL image will be warped (488) 

warp is warp array: [0] is ID, [1] are 10 X coefficients, [2] are 10 Y coefficients
source coordinates (x,y)
destination coordinates (u,v)
xWarps
u = a0*x3 + a1*x2y + a2*xy2 + a3*y3 + a4*x2 + a5*xy + a6*y2 + a7*x + a8*y + a9
yWarps
v = b0*x3 + b1*x2y + b2*xy2 + b3*y3 + b4*x2 + b5*xy + b6*y2 + b7*x + b8*y + b9

dimIndex = 1 if 3D file (TS3D), 0 if 2D file (TS2D)
zfactor: z compression
zshift is an int: z shift in nanometers (0 for a 2D file)
pxSize is the cam image pixel size in nm
*/

function TSWarp(inPath, outDir, warpA, dimIndex, zfactor, zshift, pxSize) {

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

	// coordinates header
	var xHeader = "\"x [nm]\"";
	var yHeader = "\"y [nm]\"";
	var zHeader = "\"z [nm]\"";

	// Store warp coefficients in a shorter format
	var ac = warpA[1];
	var bc = warpA[2];

	// log coeffs
	/*
	xCoefs = "X coefficients: " + ac[0];
	yCoefs = "Y coefficients: " + bc[0];
	for (w = 1; w < warpA[1].length; w++) {
		xCoefs += ", " + ac[w];
		yCoefs += ", " + bc[w];
	}
	IJ.log(xCoefs);
	IJ.log(yCoefs);
	*/
	

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

	// Read header from input file, find indexes for x, y and z coordinates
	var inLine = br.readLine();
	var inCells = inLine.split(inSep);
	var xIndex = arrayFind(inCells, xHeader);
	var yIndex = arrayFind(inCells, yHeader);
	//IJ.log("xIndex=" + xIndex + " yIndex=" + yIndex);
	if (dimIndex == 1) var zIndex = arrayFind(inCells, zHeader);
	//IJ.log("zIndex=" + zIndex);

	// Generate output name and path, open file writer
	if (dimIndex == 0) var tsExt = "TS2D";
	else tsExt = "TS3D";	
	if (outDir + inName == inPath) var outName = inName.replace(tsExt, "W_" + tsExt);
	else var outName = inName;
	var outPath = outDir + outName;
	var outFile = new File(outPath);

	// create output file and open the writer
	if (!outFile.exists()) {
		outFile.createNewFile();
	}
	IJ.log("      outName: " + outName);
	var bw = new BufferedWriter(new FileWriter(outFile));	
	
	// write header in output file
	bw.write(inLine);
	bw.newLine();

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

		ln++; 
		
		inCells = inLine.split(inSep);

		var x0 = parseFloat(inCells[xIndex] / pxSize);
		var y0 = parseFloat(inCells[yIndex] / pxSize);	

		// changed the formula to get something plausible
//		var x1 = ac[0] + ac[1]*x0 + ac[2]*ac[2]*x0*x0 + ac[3]*ac[3]*ac[3]*x0*x0*x0 + ac[4]*y0 + ac[5]*ac[5]*y0*y0 + ac[6]*ac[6]*ac[6]*y0*y0*y0 + ac[7]*ac[7]*x0*y0 + ac[8]*ac[8]*ac[8]*x0*x0*y0 + ac[9]*ac[9]*ac[9]*x0*y0*y0;
//		var y1 = bc[0] + bc[1]*x0 + bc[2]*bc[2]*x0*x0 + bc[3]*bc[3]*bc[3]*x0*x0*x0 + bc[4]*y0 + bc[5]*bc[5]*y0*y0 + bc[6]*bc[6]*bc[6]*y0*y0*y0 + bc[7]*bc[7]*x0*y0 + bc[8]*bc[8]*bc[8]*x0*x0*y0 + bc[9]*bc[9]*bc[9]*x0*y0*y0;


//		var x1 = ac[0] + ac[1]*x0 + ac[2]*ac[2]*x0*x0 + ac[3]*ac[3]*ac[3]*x0*x0*x0 + ac[4]*y0 + ac[5]*ac[5]*x0*y0 + ac[6]*ac[6]*ac[6]*x0*x0*y0 + ac[7]*ac[7]*y0*y0 + ac[8]*ac[8]*ac[8]*x0*y0*y0 + ac[9]*ac[9]*ac[9]*y0*y0*y0;
//		var y1 = bc[0] + bc[1]*x0 + bc[2]*bc[2]*x0*x0 + bc[3]*bc[3]*bc[3]*x0*x0*x0 + bc[4]*y0 + bc[5]*bc[5]*x0*y0 + bc[6]*bc[6]*bc[6]*x0*x0*y0 + bc[7]*bc[7]*y0*y0 + bc[8]*bc[8]*bc[8]*x0*y0*y0 + bc[9]*bc[9]*bc[9]*y0*y0*y0;

		var x1 = ac[0] + ac[1]*x0 + ac[2]*x0*x0 + ac[3]*x0*x0*x0 + ac[4]*y0 + ac[5]*x0*y0 + ac[6]*x0*x0*y0 + ac[7]*y0*y0 + ac[8]*x0*y0*y0 + ac[9]*y0*y0*y0;
		var y1 = bc[0] + bc[1]*x0 + bc[2]*x0*x0 + bc[3]*x0*x0*x0 + bc[4]*y0 + bc[5]*x0*y0 + bc[6]*x0*x0*y0 + bc[7]*y0*y0 + bc[8]*x0*y0*y0 + bc[9]*y0*y0*y0;

		inCells[xIndex] = (x1 * pxSize).toFixed(2);
		inCells[yIndex] = (y1 * pxSize).toFixed(2);
		
//		if (ln < 10) {
//			IJ.log("(x0,y0)=(" + x0 + "," + y0 + ")");
//			IJ.log("(x1,y1)=(" + inCells[xIndex] + "," + inCells[yIndex] + ")");
//		}
		
		if (dimIndex == 1) {
			inCells[zIndex] = ((parseFloat(inCells[zIndex]) * zfactor) + zshift).toFixed(2);
		}

		// build output line and write it into output file
		var outLine = inCells[0];
		for (var c = 1; c < inCells.length; c++) {
			outLine += sep + inCells[c];
		}
		bw.write(outLine);
		bw.newLine();
	}		
	// close reader and writer 
	br.close();
	bw.close();
}

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