/* Strahler_Analysis.bsh
 * IJ BAR: https://github.com/tferr/Scripts#scripts
 *
 * BeanShell script that performs Strahler analysis in ImageJ by repeated elimination of
 * terminal branches of topographic 2D/3D skeletons (http://fiji.sc/Strahler)
 * Tiago Ferreira, 2013-2015
 *
 * Requirements:
 * Up-to-date versions of Ignacio Arganda-Carreras Skeletonize[1] and AnalyzeSkeleton[2]
 * plugins. Both are bundled with Fiji (http://fiji.sc/) but can be otherwise downloaded
 * from http://jenkins.imagej.net/job/Stable-Fiji/ws/Fiji.app/plugins/
 * [1] http://fiji.sc/Skeletonize3D
 * [2] http://fiji.sc/AnalyzeSkeleton
 *
 * 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
 * (http://www.gnu.org/licenses/gpl.txt).
 * 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.
 */

import ij.IJ;
import ij.ImagePlus;
import ij.ImageStack;
import ij.WindowManager;
import ij.gui.GenericDialog;
import ij.gui.Roi;
import ij.plugin.ZProjector;
import ij.measure.Calibration;
import ij.measure.ResultsTable;
import ij.process.ByteProcessor;
import ij.process.ImageProcessor;
import ij.text.TextWindow;

import skeleton_analysis.AnalyzeSkeleton_;
import skeleton_analysis.Point;
import skeleton_analysis.SkeletonResult;
import Skeletonize3D_.Skeletonize3D_;

import java.awt.Rectangle;
import java.awt.image.IndexColorModel;
import java.util.ArrayList;


/* Default value for max. number of prunning cycles */
int maxPruning = 10;

/* Default option for loop detection */
int pruneChoice = AnalyzeSkeleton_.SHORTEST_BRANCH;

/* Default option for 'root-protection' ROI */
boolean protectRoot = true;

/* Default option for 'iteration-stack' output */
boolean outIS = true;

/* Default option for 'heat-map' output */
boolean outCM = true;

/* Default option for verbose mode */
boolean verbose = true;

/* Remove isolated pixels from thinned images? */
boolean erodeIsolatedPixels = true;

/* Title of results window */
String STRAHLER_TABLE = "Strahler_Table";

/* Version of this script */
String VERSION = "1.4.0 2015.01.01";

/* Flag for 'root-protective' ROI */
private static boolean validRootRoi;


/* Reminds the user to install required dependencies */
boolean validInstallation() {
	try {
		Class.forName("skeleton_analysis.AnalyzeSkeleton_");
		Class.forName("Skeletonize3D_.Skeletonize3D_");
		return true;
	} catch( ClassNotFoundException e ) {
		URL = "http://jenkins.imagej.net/job/Stable-Fiji/ws/Fiji.app/plugins/";
		AS_VRSN = "AnalyzeSkeleton_-2.0.3.jar";
		SK_VRSN = "Skeletonize3D_-1.0.1.jar";
		msg = "\n**** Strahler Analysis Error: Required file(s) not found:\n"+ e +"\n \n"
			+ "Strahler Analysis requires AnalyzeSkeleton_.jar and Skeletonize3D_.jar to be installed in\n"
			+ "the plugins/ folder. Please install the missing file(s) by double-clicking on the links below:\n \n"
			+ URL + AS_VRSN +"\n"+ URL + SK_VRSN;
		IJ.log(msg);
		return false;
	}
}

/* Checks if image fulfills analysis requirements */
boolean validImage(ImagePlus imp) {
	boolean valid = validInstallation();
	if (!valid)
		IJ.error("Strahler Analysis", "Dependencies not found.");
	valid = imp!=null && imp.getBitDepth()==8;
	if (!valid)
		IJ.error("Strahler Analysis", "An 8-bit image is required but none was found.");
	return valid;
}

/*
 * Creates the dialog prompt, retrieving the image with the original structure. While it
 * is unlikely that the iterative pruning of terminal branches will cause new loops on
 * pre-existing skeletons, offering the option to resolve loops with intensity based
 * methods remains useful specially when analyzing non-thinned grayscale images.
 */
ImagePlus getOriginalImp(String grayscaleImgChoice) {

	gd = new GenericDialog("Strahler Analysis v"+ super.VERSION);
	headerFont = new Font("SansSerif", Font.BOLD, 12);

	// Part 1. Main Options
	gd.setInsets(0, 0, 0);
	gd.addMessage("Strahler Numbering:", headerFont);
	gd.addSlider("     Max. n. of iterations:", 1, 20, super.maxPruning);
	gd.addCheckbox("Infer root end-points from rectangular ROI", super.protectRoot);
	gd.addCheckbox("Ignore single-point arbors (Isolated pixels)", super.erodeIsolatedPixels);
	if (!super.validRootRoi) {
		Checkbox roiOption = (Checkbox)gd.getCheckboxes().elementAt(0);
		roiOption.setEnabled(false);
	}

	// Part 2: Loop elimination
	gd.setInsets(25, 0, 0);
	gd.addMessage("Elimination of Skeleton Loops:", headerFont);
	gd.addChoice("Method:", AnalyzeSkeleton_.pruneCyclesModes,
					AnalyzeSkeleton_.pruneCyclesModes[super.pruneChoice]);
	int[] ids = WindowManager.getIDList();
	String[] imgTitles = new String[ids.length];
	for (int i=0; i<ids.length; ++i )
		imgTitles[i] = WindowManager.getImage(ids[i]).getTitle();
	gd.addChoice("Unsegmented image:", imgTitles, grayscaleImgChoice);

	// Part 3: Output
	gd.setInsets(25, 0, 0);
	gd.addMessage("Output Options:", headerFont);
	gd.addCheckbox("Iteration stack (Document individual iterations)", super.outIS);
	gd.addCheckbox("Heat-map image (Summarize analysis)", super.outCM);
	gd.addCheckbox("Show detailed information (Verbose mode)", super.verbose);

	gd.addHelp("http://fiji.sc/Strahler");
	gd.showDialog();
	if (gd.wasCanceled()) return null;

	super.maxPruning = (int)gd.getNextNumber();
	super.protecRoot = gd.getNextBoolean();
	super.erodeIsolatedPixels = gd.getNextBoolean();
	super.pruneChoice = gd.getNextChoiceIndex();
	imgChoice = gd.getNextChoiceIndex();
	super.outIS = gd.getNextBoolean();
	super.outCM = gd.getNextBoolean();
	super.verbose = gd.getNextBoolean();

	return WindowManager.getImage(ids[imgChoice]);
}

/* Returns the table of the specified window title */
ResultsTable getTable(String title) {
	ResultsTable rt = null;
	window = WindowManager.getFrame(title);
	if (window!=null)
		rt = ((TextWindow)window).getTextPanel().getResultsTable();
	if (rt==null)
		rt = new ResultsTable();
	return rt;
}

/* Returns the sum of elements of an int[] array */
int sum(int[] array) {
	int sum = 0;
	if (array!=null) for (int i : array) sum += i;
	return sum;
}

/* Returns the sum of elements of a double[] array */
double sum(double[] array) {
	double sum = 0;
	if (array!=null) for (double i : array) sum += i;
	return sum;
}

/*
 * Paints point positions. NB: BeanShell does not seem to suport generics:
 * paintPoints(ImageStack stack, ArrayList<Point> points, int value, String label)
 * triggers a ParseException
 */
void paintPoints(ImageStack stack, points, int value, String sliceLabel) {
	if (points!=null) {
		ImageProcessor ip = ip.createProcessor(stack.getWidth(), stack.getHeight());
		for (int j=0; j<points.size(); j++) {
			Point point = points.get(j);
			ip.putPixel(point.x, point.y, value);
		}
		stack.addSlice(sliceLabel, ip);
	}
}

/*
 * Skeletonization method that erodes the thinned structure in order to eliminate
 * isolated pixels. Thinning and pruning may give rise to single point arbors.
 * These 'debris' trees have 1 end-point but no branches or junctions. If present
 * they overestimate the total number of end-points
 */
void skeletonizeWithoutHermits(ImagePlus imp) {
	Skeletonize3D_ thin = new Skeletonize3D_();
	thin.setup("", imp);
	thin.run(null);
	if (super.erodeIsolatedPixels) {
		ImageStack stack = imp.getStack();
		for (int i=1; i<=stack.getSize(); i++)
			((ByteProcessor)stack.getProcessor(i)).erode(8, 0);
	}
}


// Retrieve analysis image and its ROI
ImagePlus srcImp = WindowManager.getCurrentImage();
if (!validImage(srcImp)) return;
String title = srcImp.getTitle();
Roi roi = srcImp.getRoi();
validRootRoi = (roi!=null && roi.getType()==Roi.RECTANGLE);

// TODO: 3D Roots are special. We need to:
// 1) Check if ROI is associated with all slices or just one
// 2) Ignore counts above/below the ROI, as needed
// 3) Extend ip operations to stack
if (validRootRoi && srcImp.getNSlices()>1) {
	IJ.error("Strahler Analysis warning", "'Root-ROI' works for 2D but not yet for 3D images.");
	validRootRoi = false;
}

// Retrieve grayscale image for intensity-based pruning of skel. loops
ImagePlus origImp = getOriginalImp(title);
if (origImp==null) return;

// Work on a skeletonized copy since we'll be modifing the image
if (roi!=null) srcImp.killRoi();
ImagePlus imp = srcImp.duplicate();
if (roi!=null) srcImp.setRoi(roi);
ImageProcessor ip = imp.getProcessor();
skeletonizeWithoutHermits(imp);

// Analyze root
ImageProcessor rootIp;
SkeletonResult rootResult;
ArrayList rootEndpointsList;

if (validRootRoi) {

	// Duplicate entire canvas. Ignore tree(s) outside ROI
	ImagePlus rootImp = imp.duplicate();
	rootIp = rootImp.getProcessor();
	rootIp.setValue(0.0);
	rootIp.fillOutside(roi);

	// Get tree properties
	AnalyzeSkeleton_ root = new AnalyzeSkeleton_();
	root.setup("", rootImp);
	rootResult = root.run(pruneChoice, false, false, origImp, true, false);
	rootEndpointsList = rootResult.getListOfEndPoints();

	// We're assuming that the ROI contains only end-point branches and slab
	// voxels but no junction points (i.e., ROI only contains disconnected
	// branches). Thus, we must remove end-points at ROI boundaries
	ListIterator it = rootEndpointsList.listIterator();
	Rectangle r = roi.getBounds();
	while (it.hasNext()) {
		Point p = it.next();
		if (p.x==r.x || p.y==r.y || p.x==r.x+r.getWidth()-1 || p.y==r.y+r.getHeight()-1)
			it.remove();
	}

}

// Initialize ResultsTable and display images. Use Z-projections
// to populate iteration stack when dealing with 3D skeletons
ResultsTable rt = getTable(STRAHLER_TABLE);
rt.setPrecision(5);

int nSlices = imp.getNSlices();
ImageStack iterationStack;
ZProjector zp;
if (outIS || outCM) {
	iterationStack = new ImageStack( imp.getWidth(), imp.getHeight() );
	if (nSlices>1) {
		zp = new ZProjector(imp);
		zp.setMethod(ZProjector.MAX_METHOD);
		zp.setStartSlice(1);
		zp.setStopSlice(nSlices);
	}
}

// Initialize AnalyzeSkeleton_
AnalyzeSkeleton_ as = new AnalyzeSkeleton_();
as.setup("", imp);

// Perform the iterative pruning
int order = 1, nEndpoints = 0, nJunctions = 0, nJunctions2 = 0;
double nEndpoints2 = Double.NaN;
ArrayList endpointsList, junctionsList;

do {

	IJ.showStatus("Retrieving measurements for order "+ order +"...");
	IJ.showProgress(order, maxPruning);
	rt.incrementCounter();

	// (Re)skeletonize image
	if (order>1) skeletonizeWithoutHermits(imp);

	// Get properties of loop-resolved tree(s)
	SkeletonResult sr = as.run(pruneChoice, false, false, origImp, true, false);
	nEndpoints = sum(sr.getEndPoints());
	nJunctions = sum(sr.getJunctions());

	// Remember initial properties
	if (order==1) {
		endpointsList = sr.getListOfEndPoints();
		junctionsList = sr.getListOfJunctionVoxels();
	}

	// Is it worth proceeding?
	if (nEndpoints==0) {
		rt.addValue("Notes", "Iteration aborted: No end-poins detected!");
		break;
	}
	if (nJunctions2==nJunctions) {
		rt.addValue("Notes", "Iteration aborted: Unsolved loop(s) detected!");
		break;
	}

	// Add current tree(s) to debug animation
	if (outIS || outCM) {
		if (nSlices>1) {
			zp.doProjection();
			ipd = zp.getProjection().getProcessor();
		} else {
			ipd = ip.duplicate();
		}
		iterationStack.addSlice("Order "+ IJ.pad(order, 2), ipd);
	}

	// Report counts
	rt.addLabel("Image", title);
	rt.addValue("Strahler order", Integer.toString(order));
	rt.addValue("# End-points", nEndpoints);
	rt.addValue("Ramification ratios", nEndpoints2/nEndpoints);
	rt.addValue("Notes", "");
	if (verbose) {
		rt.addValue("Average branch length", sum(sr.getAverageBranchLength()));
		rt.addValue("# Trees", sr.getNumOfTrees());
		rt.addValue("# Branches", sum(sr.getBranches()));
		rt.addValue("# Junctions", nJunctions);
		rt.addValue("# Triple points", sum(sr.getTriples()));
		rt.addValue("# Quadruple points", sum(sr.getQuadruples()));
	}

	// Remember main results
	nEndpoints2 = (double)nEndpoints;
	nJunctions2 = nJunctions;

	// Eliminate end-points
	as.run(pruneChoice, true, false, origImp, true, false, roi);

} while (order++ <= maxPruning && nJunctions > 0);


if (validRootRoi) {

	// Check if ROI contains unexpected structures
	int junctions = sum(rootResult.getJunctions());
	String msg = (junctions > 0) ? "Warning: ROI contains ramified root(s)"
				: "Root-branches inferred from ROI";

	// Append root properties to table and display results
	rt.incrementCounter();
	rt.addLabel("Image", title);
	rt.addValue("Strahler order", Integer.toString(order)+" (Root)");
	rt.addValue("# End-points", rootEndpointsList.size());
	rt.addValue("Ramification ratios", nEndpoints2/rootEndpointsList.size());
	rt.addValue("Notes", msg);
	if (verbose) {
		rt.addValue("Average branch length", sum(rootResult.getAverageBranchLength()));
		rt.addValue("# Trees", rootResult.getNumOfTrees());
		rt.addValue("# Branches", sum(rootResult.getBranches()));
		rt.addValue("# Junctions", junctions);
		rt.addValue("# Triple points", sum(rootResult.getTriples()));
		rt.addValue("# Quadruple points", sum(rootResult.getQuadruples()));
	}
}
rt.show(STRAHLER_TABLE);

// Create results image(s)
if (iterationStack.getSize()>1 && (outIS || outCM)) {
	ImagePlus imp2 = new ImagePlus("StrahlerAnimation_"+ title, iterationStack);
	Calibration cal = srcImp.getCalibration();
	if (outIS) {
		if (validRootRoi) {
			iterationStack.addSlice("Root", rootIp);
			paintPoints(iterationStack, rootEndpointsList, 255, "Root end-points");
			imp2.setRoi(roi);
		}
		paintPoints(iterationStack, endpointsList, 255, "End-points");
		paintPoints(iterationStack, junctionsList, 255, "Junction-points");
		imp2.show();
	}
	if (outCM) {
		zp = new ZProjector(imp2);
		zp.setMethod(ZProjector.SUM_METHOD);
		zp.setStartSlice(1);
		zp.setStopSlice(order-1);
		zp.doProjection();
		ImageProcessor ip3 = zp.getProjection().getProcessor();
		ip3.multiply(1/255.0);
		ip3.setMinAndMax(0, order-1);
		ImagePlus imp3 = new ImagePlus("StrahlerMask_"+ title, ip3);
		try {
			Class.forName("Sholl_Utils");
			ip3.setColorModel(Sholl_Utils.matlabJetColorMap(0));
		} catch (ClassNotFoundException) {
			IJ.run(imp3, "Fire", "");
		}
		IJ.run(imp3, "Calibration Bar...", " overlay");
		if (validRootRoi) imp3.setRoi(roi);
		imp3.show();
	}
}

IJ.showProgress(0, 0);
IJ.showTime(imp, imp.getStartTime(), "Strahler Analysis concluded... ");