/* 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-2016
 *
 * 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.DialogListener;
import ij.gui.GenericDialog;
import ij.gui.ImageCanvas;
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 sc.fiji.analyzeSkeleton.AnalyzeSkeleton_;
import sc.fiji.analyzeSkeleton.Point;
import sc.fiji.analyzeSkeleton.SkeletonResult;
import Skeletonize3D_.Skeletonize3D_;

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


// Before starting we must remind non-Fiji users to install required dependencies
try {
	Class.forName("sc.fiji.analyzeSkeleton.AnalyzeSkeleton_");
	Class.forName("Skeletonize3D_.Skeletonize3D_");
} catch( ClassNotFoundException e ) {
	URL = "http://jenkins.imagej.net/job/Stable-Fiji/ws/Fiji.app/plugins/";
	AS_VRSN = "AnalyzeSkeleton_-3.0.0.jar";
	SK_VRSN = "Skeletonize3D_-1.0.2.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.handleException(e);
	IJ.log(msg);
	IJ.error("Strahler Analysis Error",
		"Dependencies not found!\nCheck the Log window for installation details.");
	return;
}

/* 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 = false;

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

/* Default option for tabular option */
boolean tabular = false;

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

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

/* Title of detailed results window */
String VERBOSE_TABLE = "Strahler_Iteration_Log";

/* Version of the program */
String VERSION = "1.5.1 2016.01.19";

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


/* Checks if image fulfills analysis requirements */
boolean validImage(ImagePlus imp) {
	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) {

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

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

	// 8-bit grayscale is the only image type recognized by AnalyzeSkeleton_,
	// so we'll provide the user with a pre-filtered list of valid choices
	ArrayList validIds = new ArrayList();
	ArrayList validTitles = new ArrayList();
	int[] ids = WindowManager.getIDList();
	for (int i=0; i<ids.length; ++i ) {
		ImagePlus imp = WindowManager.getImage(ids[i]);
		if (imp.getBitDepth()==8) { //TODO: ignore composites?
			validIds.add(ids[i]);
			validTitles.add(imp.getTitle());
		}
	}
	gd.addChoice("8-bit grayscale image:", validTitles.toArray(new String[validTitles.size()]),
		grayscaleImgChoice);

	// Part 3: Output
	gd.setInsets(25, 0, 0);
	gd.addMessage("Output Options:", headerFont);
	gd.addCheckbox("Display_iteration stack", super.outIS);
	gd.addCheckbox("Show detailed information", super.verbose);
	gd.addCheckbox("Tabular data only (no image output)", super.tabular);

	gd.addHelp("http://fiji.sc/Strahler");
	gd.addDialogListener(this);
	dialogItemChanged(gd, null);
	gd.showDialog();

	return (gd.wasCanceled()) ? null : WindowManager.getImage(validIds.get(imgChoice));
}

/* Retrive dialog options using the DialogListener interface */
boolean dialogItemChanged(GenericDialog gd, java.awt.AWTEvent e) {

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

	// Enable/Disable key components of GenericDialog
	Choice imgChoice = (Choice)gd.getChoices().elementAt(1);
	imgChoice.setEnabled((super.pruneChoice>1) ? true : false);
	Vector checkboxes = gd.getCheckboxes();
	Checkbox roiOption = (Checkbox)checkboxes.elementAt(0);
	roiOption.setEnabled(super.validRootRoi);
	Checkbox stackOption = (Checkbox)checkboxes.elementAt(2);
	stackOption.setEnabled(!super.tabular);
	return !gd.wasCanceled();

}

/* Returns the table of the specified window title setting common properties */
ResultsTable getTable(String title) {
	ResultsTable rt = null;
	Window window = WindowManager.getWindow(title);
	if (window!=null)
		rt = ((TextWindow)window).getTextPanel().getResultsTable();
	if (rt==null)
		rt = new ResultsTable();
	rt.setPrecision(5);
	rt.setNaNEmptyCells(true);
	rt.showRowNumbers(false);
	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;
}

/* Returns the average of an int[]/double[] array */
double average(array) {
	return sum(array) / array.length;
}

/*
 * 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);
	}
}

/* Clears point positions */
void clearPoints(ImageProcessor ip, points) {
	if (points!=null) {
		for (int j=0; j<points.size(); j++) {
			Point point = points.get(j);
			ip.putPixel(point.x, point.y, 0);
		}
	}
}

/*
 * 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);
	}
}

/* Runs ij.plugin.CalibrationBar on the specified image using a suitable scale */
void addCalibrationBar(ImagePlus imp, int nLabels) {
	ImageCanvas ic = imp.getCanvas();
	double zoom = 1.0;
	double mag = (ic!=null) ? ic.getMagnification() : 1.0;
	if (zoom<=1 && mag<1)
		zoom = (double) 1.0/mag;
	IJ.run(imp, "Calibration Bar...",
		"fill=Black label=White number="+ nLabels +" zoom="+ zoom +" overlay");
}

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

// Initialize ResultsTable: main and detailed info
ResultsTable rt = getTable(STRAHLER_TABLE);
ResultsTable logrt = getTable(VERBOSE_TABLE);

// Analyze root
ImagePlus rootImp;
ImageProcessor rootIp;
SkeletonResult rootResult;
ArrayList rootEndpointsList;
int nRootEndpoints = 0, nRootJunctions = 0;

if (validRootRoi && verbose) {

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

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

	// We assume ROI contains only end-point branches, slab voxels and
	// no junction points. We'll thus remove end-points at ROI boundaries
	nRootJunctions = sum(rootResult.getJunctions());
	rootEndpointsList = rootResult.getListOfEndPoints();
	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();
	}
	rootResult.setListOfEndPoints(rootEndpointsList);
	nRootEndpoints = rootEndpointsList.size();

}

// Initialize display images. Use Z-projections to populate
// iteration stack when dealing with 3D skeletons
int nSlices = imp.getNSlices();
ZProjector zp;

ImageStack 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
boolean aborted = false;
int order = 1, nEndpoints = 0, nJunctions = 0, nJunctions2 = 0;
ArrayList endpointsList, junctionsList;
String errorMsg = "";

do {

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

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

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

		// Do not include root in 1st order calculations
		nEndpoints -= nRootEndpoints;
		nJunctions -= nRootJunctions;
	}

	// Is it worth proceeding?
	if (nEndpoints==0 || nJunctions2==nJunctions) {
		aborted = true;
		errorMsg = "Error! Iteration "+ order +" aborted: ";
		errorMsg += (nEndpoints==0) ? "No end-poins found" : "Unsolved loop(s) detected";
		break;
	}

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

	// Report properties of pruned structures
	if (verbose) {
		logrt.incrementCounter();
		logrt.addLabel("Image", title);
		logrt.addValue("Structure", "Skel. at iteration "+ Integer.toString(order));
		logrt.addValue("Notes", errorMsg);
		logrt.addValue("# Trees", sr.getNumOfTrees());
		logrt.addValue("# Branches", sum(sr.getBranches()));
		logrt.addValue("# End-points", nEndpoints);
		logrt.addValue("# Junctions", nJunctions);
		logrt.addValue("# Triple points", sum(sr.getTriples()));
		logrt.addValue("# Quadruple points", sum(sr.getQuadruples()));
		logrt.addValue("Average branch length", average(sr.getAverageBranchLength()));
	}

	// Remember main results
	nJunctions2 = nJunctions;

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

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

// Set counter to the de facto order
order -= 1;

// Append root properties to log table
if (validRootRoi && verbose) {

	// Check if ROI contains unexpected structures
	String msg = (nRootJunctions > 0) ? "Warning: ROI contains ramified root(s)"
				: "Root-branches inferred from ROI";
	logrt.incrementCounter();
	logrt.addLabel("Image", title);
	logrt.addValue("Structure", "Root");
	logrt.addValue("Notes", msg);
	logrt.addValue("# Trees", rootResult.getNumOfTrees());
	logrt.addValue("# Branches", sum(rootResult.getBranches()));
	logrt.addValue("# End-points", nRootEndpoints);
	logrt.addValue("# Junctions", nRootJunctions);
	logrt.addValue("# Triple points", sum(rootResult.getTriples()));
	logrt.addValue("# Quadruple points", sum(rootResult.getQuadruples()));
	logrt.addValue("Average branch length", average(rootResult.getAverageBranchLength()));

}

// Safety check
if (iterationStack==null || iterationStack.getSize()<1) { //TODO: Make error more informative
	IJ.error("Strahler Analysis Error", "Could not complete analysis!\n"
		+ "Enable verbose mode and check "+ VERBOSE_TABLE + " for details.");
	return;
}

// Create iteration stack
Calibration cal = srcImp.getCalibration();
ImagePlus imp2 = new ImagePlus("StrahlerIteration_"+ title, iterationStack);
imp2.setCalibration(cal);
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");
}

// Generate Strahler mask
zp = new ZProjector(imp2);
zp.setMethod(ZProjector.SUM_METHOD);
zp.setStartSlice(1);
zp.setStopSlice(order);
zp.doProjection();
ImageProcessor ip3 = zp.getProjection().getProcessor().convertToShortProcessor(false);
clearPoints(ip3, junctionsList); // disconnect branches
ip3.multiply(1/255.0); // map intensities to Strahler orders
ImagePlus imp3 = new ImagePlus("StrahlerMask_"+ title, ip3);
imp3.setCalibration(cal);

// Measure segmented orders
double prevNbranches = Double.NaN;
for (int i=1; i<=order; i++) {

	// Segment branches by order
	ImagePlus maskImp = imp3.duplicate(); // Calibration will be retained
	IJ.setThreshold(maskImp, i, i); //TODO: Use ImagePlus/ImageProcessor API
	IJ.run(maskImp, "Convert to Mask", "");

	// Analyze segmented order
	AnalyzeSkeleton_ maskAs = new AnalyzeSkeleton_();
	maskAs.setup("", maskImp);
	SkeletonResult maskSr = maskAs.run(pruneChoice, false, false, origImp, true, false);
	maskImp.flush();

	// Since all branches are disconnected at this stage, the n. of branches is
	// the same as the # the trees unless zero-branches trees exist, i.e., trees
	// with no slab voxels (defined by just an end-point). We will ignore those
	// trees if the user requested it
	int nBranches = (erodeIsolatedPixels) ? sum(maskSr.getBranches()) : maskSr.getNumOfTrees();

	// Log measurements
	rt.incrementCounter();
	rt.addLabel("Image", title);
	rt.addValue("Strahler Order", i);
	rt.addValue("# Branches", nBranches);
	rt.addValue("Ramification ratios", prevNbranches/nBranches);
	rt.addValue("Average branch length", average(maskSr.getAverageBranchLength()));
	rt.addValue("Unit", cal.getUnit());
	String noteMsg = "";
	if (i==1) {
		noteMsg = (erodeIsolatedPixels) ? "Ignoring" : "Including";
		noteMsg += " single-point arbors...";
	}
	rt.addValue("Notes", noteMsg);

	// Remember results for previous order
	prevNbranches = (double)nBranches;

}

// Append any errors to last row
rt.addValue("Notes", errorMsg);

// Display outputs
if (!tabular) {
	if (outIS)
		imp2.show();
	ip3.setMinAndMax(0, order);
	try {
		Class.forName("Sholl_Utils");
		ip3.setColorModel(Sholl_Utils.matlabJetColorMap(0));
	} catch (ClassNotFoundException) {
		IJ.run(imp3, "Fire", "");
	}
	if (validRootRoi)
		imp3.setRoi(roi);
	imp3.show();
	addCalibrationBar(imp3, Math.min(order, 5));
}
if (verbose)
	logrt.show(VERBOSE_TABLE);
rt.show(STRAHLER_TABLE);

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