/** Align_3D / Bone_3D_Moments
 *tool to calculate centroid and principal axes 
 *of a thresholded stack; assumes a 16-bit CT scan 
 *of a bone in air - default thresholds are 0 and 4000 HU
 *bone density assumed to be 1.8 g/cm^3 
 *Eigen decomposition performed with the Jama package
 *http://math.nist.gov/javanumerics/jama/
 *Outputs stack data aligned to principal axes
 *
 *This program is free software: you can redistribute it and/or modify
 *it under the terms of the GNU General Public License as published by
 *the Free Software Foundation, either version 3 of the License, or
 *(at your option) any later version.
 *
 *This program is distributed in the hope that it will be useful,
 *but WITHOUT ANY WARRANTY; without even the implied warranty of
 *MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *GNU General Public License for more details.
 *
 *You should have received a copy of the GNU General Public License
 *along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *@author Michael Doube
 *@version 0.2
 */

import java.awt.Rectangle;
import ij.IJ;
import ij.ImagePlus;
import ij.ImageStack;
import ij.process.ImageProcessor;
import ij.plugin.filter.PlugInFilter;
import ij.measure.Calibration;
import ij.measure.ResultsTable;
import ij.gui.*;
import Jama.*;

public class Bone_3D_Moments implements PlugInFilter {
	ImagePlus imp;
	protected ImageStack stack;
	
	public int setup(String arg, ImagePlus imp) {
	    stack = imp.getStack();
	    this.imp = imp;
	    return DOES_16 + STACK_REQUIRED + SUPPORTS_MASKING;
	}
	
	public void run(ImageProcessor ip) {
		ResultsTable rt3D = ResultsTable.getResultsTable();
		short minT = 0;			//minimum bone value in HU 
		short maxT = 4000;		//maximum bone value in HU
		Calibration cal = imp.getCalibration();
	    double vW = cal.pixelWidth;
	    double vH = cal.pixelHeight;
	    double vD = cal.pixelDepth;
	    String units = cal.getUnits();
	    double[] coeff = cal.getCoefficients();
	    float[] CTable = cal.getCTable();
	    String valueUnit = cal.getValueUnit();
	    if (cal.calibrated()) {
	    	//convert HU limits to pixel values
	    	IJ.log("Image is calibrated, using "+minT+" and "+maxT+" "+valueUnit+" as bone cutoffs");
	    	minT = (short)Math.round(cal.getRawValue(minT));
	    	maxT = (short)Math.round(cal.getRawValue(maxT));
	    	IJ.log("Vox Width: "+vW+"; Vox Height: "+vH+" Vox Depth: "+vD+" "+units);
	    	IJ.log("Calibration coefficients:"+coeff[0]+","+coeff[1]);
	    	IJ.log("minT = "+minT+", maxT = "+maxT);
	    }
	    else {
	    	IJ.log("Image is not calibrated, using user-determined threshold");
    		IJ.run("Threshold...");
    		new WaitForUserDialog("This image is not density calibrated.\nSet the threshold, then click OK.").show();
    		minT = (short)ip.getMinThreshold();
    		maxT = (short)ip.getMaxThreshold();
    	}
	    String title = imp.getTitle();
	    int startSlice = 1;
	    int endSlice = stack.getSize();
        GenericDialog gd = new GenericDialog("Limit Slices");
        gd.addNumericField("Start Slice:",startSlice,0);
        gd.addNumericField("End Slice:",endSlice,0);
        gd.showDialog();
        if (gd.wasCanceled()) {
            IJ.error("PlugIn canceled!");
            return;
        }
        startSlice = (int)gd.getNextNumber();
        endSlice = (int)gd.getNextNumber();
		int al = stack.getSize() + 1;
		//2D centroids
		double[] xc; double[] yc;
		xc = new double[al]; yc = new double[al];
		//3D centroids
		double xcent = 0; double ycent = 0; double zcent = 0;
		//pixel counters
		double cstack = 0;
		Rectangle r = ip.getRoi();
	    int offset, i;
	    int w = stack.getWidth();
	    int h = stack.getHeight();
	    int sliceSize = w*h;
	    boolean[] emptySlices;
	    emptySlices = new boolean[al];
	    short[] slicePixels = new short[sliceSize];
		double[] cslice;
    	cslice = new double[al];
    	int sl = w*h*al;
    	int maxIndex = sl-1;
    	short[] sourceWorkArray = new short[sl];
	    for (int n = 0; n<sourceWorkArray.length; n++) sourceWorkArray[n] = 0;
    	short[] targetWorkArray = new short[sl];
	    for (int s=startSlice;s<=endSlice;s++) {
	    	IJ.showStatus("Calculating centroid...");
	    	IJ.showProgress(s, endSlice);
	    	double sumx = 0; double sumy = 0;
	    	cslice[s] = 0;
	    	slicePixels = (short[])stack.getPixels(s);
	    	for (int y=r.y; y<(r.y+r.height); y++) {
		        offset = y*w;
		        for (int x=r.x; x<(r.x+r.width); x++) {
		            i = offset + x;
		            sourceWorkArray[s*sliceSize+y*w+x] = slicePixels[i];
		            if (slicePixels[i] >= minT && slicePixels[i] <= maxT){
		            	cslice[s]++;
		            	sumx += (double)x;
		            	sumy += (double)y;
		            }
		        }
		    }
		    if (cslice[s] > 0){
		    	xc[s] = sumx/cslice[s];
			    yc[s] = sumy/cslice[s];
			    xcent += xc[s]*cslice[s];
			    ycent += yc[s]*cslice[s];
			    zcent += cslice[s]*s;
			    cstack += cslice[s];
			    emptySlices[s] = false;
		    } else {
		    	emptySlices[s] = true;
		    }
	    }
	    double centX3D; double centY3D; double centZ3D;
	    int centXV; int centYV; int centZV;
	    if (cstack > 0){
	    	centX3D = xcent*vW/cstack; centY3D = ycent*vH/cstack; centZ3D = zcent*vD/cstack; //centroid in real units
	    	centXV = (int)Math.round(xcent/cstack); centYV = (int)Math.round(ycent/cstack); centZV = (int)Math.round(zcent/cstack); //centroid in voxels 
	    } else {
	    	centX3D = -1; centY3D = -1; centZ3D = -1;
	    	IJ.error("Empty Stack","No voxels are available for calculation.\nCheck your ROI and threshold.");
	    	return;
	    }
	    //END OF CENTROID CALCULATION
		//START OF 3D MOMENT CALCULATIONS
	    //Our CT scans are not quantitative as they contain sharpening artefact
	    //so density (g.cm^-3) is left out of this calculation and estimated at the end
		double Icxx = 0; double Icyy = 0; double Iczz = 0;
		double Icxy = 0; double Icxz = 0; double Icyz = 0;
		for (int z=startSlice;z<=endSlice;z++) {
			IJ.showStatus("Calculating inertia tensor...");
			IJ.showProgress(z, endSlice);
	    	if (!emptySlices[z]){
	    		slicePixels = (short[])stack.getPixels(z);
	    		for (int y=r.y; y<(r.y+r.height); y++) {
	    			offset = y*w;
	    			for (int x=r.x; x<(r.x+r.width); x++) {
	    				i = offset + x;
	    				if (slicePixels[i] >= minT && slicePixels[i] <= maxT){
	    					//moments around the 3 orthogonal axes which are parallel to the 
	    					//reference frame's axes and whose intersection is the 3D centroid
	    					Icxx += (y*vH-centY3D)*(y*vH-centY3D)+(z*vD-centZ3D)*(z*vD-centZ3D);
	    					Icyy += (x*vW-centX3D)*(x*vW-centX3D)+(z*vD-centZ3D)*(z*vD-centZ3D);
	    					Iczz += (y*vH-centY3D)*(y*vH-centY3D)+(x*vW-centX3D)*(x*vW-centX3D);
	    					Icxy += (x*vW-centX3D)*(y*vH-centY3D);
	    					Icxz += (x*vW-centX3D)*(z*vD-centZ3D);
	    					Icyz += (y*vH-centY3D)*(z*vD-centZ3D);
	    				}
	    			}
	    		}
	    	}
		}
		//we have to multiply each element of the sum by the mass of the element
		//which is vW*vH*vD * rho.  Since all voxels are the same, we can 
		//multiply the sum by voxel volume
		double voxelVolume = vW*vH*vD;
		Icxx *= voxelVolume;
		Icyy *= voxelVolume;
		Iczz *= voxelVolume;
		Icxy *= voxelVolume;
		Icxz *= voxelVolume;
		Icyz *= voxelVolume;
		//add a term for the moment of each voxel around its COM, parallel to each axis
		double shapeVolume = cstack*voxelVolume;
	    Icxx += shapeVolume*(vH*vH + vD*vD) / 12;
	    Icyy += shapeVolume*(vW*vW + vD*vD) / 12;
	    Iczz += shapeVolume*(vH*vH + vW*vW) / 12;
	    //then we multiply by density later to get the mass term
	    
	    //create the intertia tensor matrix
	    double[][] inertiaTensor = new double[3][3];
	    inertiaTensor[0][0] = Icxx;
	    inertiaTensor[1][1] = Icyy;
	    inertiaTensor[2][2] = Iczz;
	    inertiaTensor[0][1] = -Icxy;
	    inertiaTensor[0][2] = -Icxz;
	    inertiaTensor[1][0] = -Icxy;
	    inertiaTensor[1][2] = -Icyz;
	    inertiaTensor[2][0] = -Icxz;
	    inertiaTensor[2][1] = -Icyz;
	    Matrix inertiaTensorMatrix = new Matrix(inertiaTensor);
	    
	    //do the Eigenvalue decomposition
	    EigenvalueDecomposition E = new EigenvalueDecomposition(inertiaTensorMatrix);
        Matrix eVal = E.getD();
        Matrix eVec = E.getV();
        double[][] eigenValues = eVal.getArrayCopy();
        double[][] eigenVectors = eVec.getArrayCopy();
        //check orientation of eigenvectors and correct them
        if (eigenVectors[2][0] < 0){
//        	eVec = eVec.times(-1);  //this changes coordinate frame from right-hand- to left-hand-rule
        	double[][] eVecA = eVec.getArray();
        	//rotate 180 deg
        	for (int row = 0; row < 3; row++){
        		for (int col = 0; col < 2; col++){
        			eVecA[row][col] *= -1;
        		}
        	}
           	eigenVectors = eVec.getArrayCopy();
        }
        Matrix eVecInv = eVec.inverse();
        double[][] eigenVectorsInverse = eVecInv.getArrayCopy();
        
        //sometimes the eigen vector is pointing 180 deg in the wrong direction - why? HOw to fix?
        IJ.log("\nEigenvector matrix");
	    for (int j = 0; j<3; j++){
	    	IJ.log("||"+eigenVectors[j][0]+" | "+eigenVectors[j][1]+" | "+eigenVectors[j][2]+"||");
	    }
	    
	    IJ.log("\nInverse Eigenvector matrix");
	    for (int j = 0; j<3; j++){
	    	IJ.log("||"+eigenVectorsInverse[j][0]+" | "+eigenVectorsInverse[j][1]+" | "+eigenVectorsInverse[j][2]+"||");
	    }
	    //cortical bone apparent density (material density * volume fraction) from Mow & Huiskes (2005) p.140 
	    //using 1.8 g.cm^-3: 1mm^3 = 0.0018 g = 0.0000018 kg = 1.8*10^-6 kg; 1mm^2 = 10^-6 m^2
	    //conversion coefficient from mm^5 to kg.m^2 = 1.8*10^-12
	    double cc = 1.8*Math.pow(10, -12);
	    rt3D.incrementCounter();
		rt3D.addLabel("Label", imp.getTitle());
		rt3D.addValue("Xc ("+units+")", centX3D);
		rt3D.addValue("Yc ("+units+")", centY3D);
		rt3D.addValue("Zc ("+units+")", centZ3D);
		rt3D.addValue("Vol (mm^3)", shapeVolume);
		rt3D.addValue("Icxx (kg.m^2)", Icxx*cc);
		rt3D.addValue("Icyy (kg.m^2)", Icyy*cc);
		rt3D.addValue("Iczz (kg.m^2)", Iczz*cc);
		rt3D.addValue("Icxy (kg.m^2)", Icxy*cc);
		rt3D.addValue("Icxz (kg.m^2)", Icxz*cc);
		rt3D.addValue("Icyz (kg.m^2)", Icyz*cc);
		rt3D.addValue("I1 (kg.m^2)", eigenValues[2][2]*cc);
		rt3D.addValue("I2 (kg.m^2)", eigenValues[1][1]*cc);
		rt3D.addValue("I3 (kg.m^2)", eigenValues[0][0]*cc);
		rt3D.show("Results");
	    //for each voxel in the target stack, find the corresponding source voxel
		for (int z=1; z<al; z++){
			IJ.showStatus("Aligning image stack...");
			IJ.showProgress(z, al);
			if (!emptySlices[z]){
				for (int y=0; y<h; y++){
					for (int x=0; x<w; x++){
						double xD = x*vW - centX3D;
						double yD = y*vH - centY3D;
						double zD = z*vD - centZ3D;
						double xAlign = xD*eigenVectorsInverse[2][0] + yD*eigenVectorsInverse[1][0] + zD*eigenVectorsInverse[0][0] + centX3D;
						double yAlign = xD*eigenVectorsInverse[2][1] + yD*eigenVectorsInverse[1][1] + zD*eigenVectorsInverse[0][1] + centY3D;
						double zAlign = xD*eigenVectorsInverse[2][2] + yD*eigenVectorsInverse[1][2] + zD*eigenVectorsInverse[0][2] + centZ3D;
						//possibility to do some voxel interpolation instead
						//of just rounding in next 3 lines
						int xA = (int)Math.round(xAlign/vW);
						int yA = (int)Math.round(yAlign/vH);
						int zA = (int)Math.round(zAlign/vD);
						int k = zA*sliceSize+yA*w+xA;
						//possibility of resizing the target stack depending on the
						//dimensions of the aligned image rather than just clipping
						k = Math.max(k, 0);
						k = Math.min(k, maxIndex); //clip k to limits of sourceWorkArray's index
						targetWorkArray[z*sliceSize+y*w+x] = sourceWorkArray[k];
					}
				}
			}
		}
		//draw axes on stack
		int centPos = centYV*w+centXV;
		for (int z=1; z<al; z++){
			int zOffset = z*sliceSize; 
			targetWorkArray[zOffset+centPos] = maxT;
			if (z == centZV){
				//y axis
				for (int y=0; y<h; y++) targetWorkArray[zOffset+y*w+centXV] = maxT;
				//x axis
				for (int x=0; x<w; x++)	targetWorkArray[zOffset+centYV*w+x] = maxT;
			}
		}
		ImageStack target = new ImageStack(w,h);
		for (int z = 1; z < al; z++){
			short[] sliceArray = new short[sliceSize];
	    	System.arraycopy(targetWorkArray, z*sliceSize, sliceArray, 0, sliceSize);
	    	target.addSlice("slice "+z, sliceArray);
	    }
		ImagePlus impTarget = new ImagePlus(title+"_aligned", target);
		Calibration set = impTarget.getCalibration();
		set.pixelWidth = vW;
		set.pixelHeight = vH;
		set.pixelDepth = vD;
		set.setUnit(units);
		set.setCTable(CTable, valueUnit);
		impTarget.show();
		IJ.setMinAndMax(minT+coeff[0], maxT+coeff[0]);
		impTarget.updateAndDraw();
	}
}