/* 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
 * Tiago Ferreira, v1.3.2 2014.11.18
 *
 * Requirements:
 * Ignacio Arganda-Carreras Skeletonize (http://fiji.sc/Skeletonize3D) and AnalyzeSkeleton
 * (http://fiji.sc/AnalyzeSkeleton) plugins, both bundled with Fiji (http://fiji.sc/)
 *
 * Installation:
 * Save this file in the plugins/ folder using the 'Plugins>Install...' command.
 *
 * Resources:
 * http://jenkins.imagej.net/job/Fiji-javadoc/javadoc/Skeletonize3D_/Skeletonize3D_.html
 * http://jenkins.imagej.net/job/Fiji-javadoc/javadoc/skeleton_analysis/package-summary.html
 *
 * 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.Overlay;
import ij.gui.Roi;
import ij.plugin.ImageCalculator;
import ij.plugin.ZProjector;
import ij.plugin.filter.ParticleAnalyzer;
import ij.measure.Calibration;
import ij.measure.ResultsTable;
import ij.process.ImageProcessor;
import ij.text.TextWindow;

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


/* Definitions */
String[] OUT_CHOICES = {"None", "Iteration Stack (IS)", "Color Map (CM)", "Both IS & CM"};
int OUT_NONE = 0; int OUT_IS = 1; int OUT_CM = 2; int OUT_ISCM = 3;
boolean protectRoot;

/* Default settings */
int maxPruning = 10;        // Max. number of prunning cycles
int outChoice = OUT_CM;     // Default choice for output image

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


/* 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.2.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);
        lw = WindowManager.getFrame("Log");
        if (lw!=null) lw.setSize(645, 270);
        return false;
    }
}

/*
 * 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(ImagePlus currentImp, String currentTitle) {

    ImagePlus origImg = currentImp;
    gd = new GenericDialog("Strahler Analysis of "+ currentTitle);

    // Part 1. Main Options
    createDialogHeader(gd, "I. Strahler Numbering:", true);
    gd.addSlider("Max. n. of iterations:", 1, 20, maxPruning);
    gd.addChoice("Output image(s):", OUT_CHOICES , OUT_CHOICES[outChoice]);

    // Part 2: Root protection
    createDialogHeader(gd, "II. Non-radial Arbors:", false);
    String protectMsg;
    if (protectRoot) {
        gd.addCheckbox("Protect ROI from iterative pruning", protectRoot);
        protectMsg = "ROI should contain a single branch corresponding\nto the root node.";
    } else
        protectMsg = "No area ROI was found. All segments will be exposed\n"
                   + "to the iterative elimination of terminal branches.";

    gd.setInsets(0, 20, 0);
    gd.addMessage(protectMsg);

    // Part 3: Loop elimination
    createDialogHeader(gd, "III. Elimination of Skeleton Loops:", false);
    int[] ids = WindowManager.getIDList();
    boolean singleImage = (ids==null || ids.length<2);

    if (singleImage) { // Only skeleton image is open: intensity based pruning is not available

        String[] S_CHOICES = { AnalyzeSkeleton_.pruneCyclesModes[AnalyzeSkeleton_.NONE],
                    AnalyzeSkeleton_.pruneCyclesModes[AnalyzeSkeleton_.SHORTEST_BRANCH]};
        gd.addChoice("Method:", S_CHOICES, S_CHOICES[pruneChoice]);

    } else { // Non-thinned image may be available: Allow intensity based pruning

        gd.addChoice("Method:", AnalyzeSkeleton_.pruneCyclesModes,
                    AnalyzeSkeleton_.pruneCyclesModes[pruneChoice]);
        String[] imgTitles = new String[ids.length];
        for (int i=0; i<ids.length; ++i )
            imgTitles[i] = (WindowManager.getImage(ids[i])).getTitle();
        gd.addChoice("Original grayscale image:", imgTitles, currentTitle);

    }

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

    maxPruning = (int)gd.getNextNumber();
    outChoice = gd.getNextChoiceIndex();
    if (protectRoot) protectRoot = gd.getNextBoolean();
    pruneChoice = gd.getNextChoiceIndex();
    if (!singleImage) {
        imgChoice = gd.getNextChoiceIndex();
        if (pruneChoice==AnalyzeSkeleton_.LOWEST_INTENSITY_VOXEL || pruneChoice==AnalyzeSkeleton_.LOWEST_INTENSITY_BRANCH)
            origImg = WindowManager.getImage(ids[imgChoice]);
    }

    return origImg;
}

/* Outputs measurements to a dedicated table */
void printResults(order, tips, ratios) {
    strahlerTable = "Strahler_Table";
    window = WindowManager.getFrame(strahlerTable);
    if (window == null)
        rt = new ResultsTable();
    else
        rt = ((TextWindow) window).getTextPanel().getResultsTable();
    rt.setPrecision(0);
    rt.incrementCounter();
    rt.addLabel("Image", title);
    rt.addValue("Strahler order", order);
    rt.addValue("# End-point branches", tips);
    rt.addValue("Ramification ratios", IJ.d2s(ratios,5));
    rt.show(strahlerTable);
}

/* Adds a section header to a GenericDialog */
void createDialogHeader(GenericDialog gd, String text, boolean supressAboveSpace) {
     gd.setInsets(supressAboveSpace?0:12, 0, 0);
     gd.addMessage(text, new Font("SansSerif", Font.BOLD, 12));
}

/* Adds "Root ROI" to the image overlay */
void addRootToOverlay(ImagePlus imp, Roi roi, boolean label) {
    Overlay overlay = imp.getOverlay();
    if (overlay==null) overlay = new Overlay();
    roi.setName("Root");
    overlay.add(roi);
    if (label) {
        overlay.setLabelColor(roi.getColor());
        overlay.drawNames(true);
    }
    imp.setOverlay(overlay);
}

/*
 * Paints junction points. NB: BeanShell does not seem to suport generics. This triggers
 * a ParseException: paintBranchPoints(ImageProcessor ip, ArrayList<Point> points, int value)
 */
void paintBranchPoints(ImageProcessor ip, points, int value) {
    if (points!=null) {
        for (int j=0; j<points.size(); j++) {
            Point point = points.get(j);
            ip.putPixel(point.x, point.y, value);
        }
    }
}


// Startup checks
if (!validInstallation())
    { IJ.error("Missing files? See Log window for details."); return; }
ImagePlus srcImp = WindowManager.getCurrentImage();
if (srcImp==null)
    { IJ.error("There are no images open."); return; }
if (srcImp.getBitDepth()!=8)
    { IJ.error("Cannot process "+ srcImp.getTitle() +".\nOnly 8-bit images can be processed."); return; }

// Root protection?
Roi roi = srcImp.getRoi();
boolean protectRoot = roi!=null && roi.isArea();

// Prompt for options and retrieve image for intensity-based pruning of skel. loops
ImagePlus origImp = getOriginalImp(srcImp, title = srcImp.getTitle() );
if (origImp==null) return;

// Work on a copy, since the original image will be modified
if (roi!=null) srcImp.killRoi();
ImagePlus imp = srcImp.duplicate();
ImageProcessor ip = imp.getProcessor();
if (roi!=null) srcImp.setRoi(roi);

// Skeletonize in case image is not a skeleton
if (roi!=null) imp.killRoi();
Skeletonize3D_ thin = new Skeletonize3D_();
thin.setup("", imp);
thin.run(ip);

// Extract coordinates of all junction pixels
initialSkel = new AnalyzeSkeleton_();
initialSkel.setup("", imp);
initialResult = initialSkel.run(pruneChoice, false, false, origImp, true, false);
junctionsList = initialResult.getListOfJunctionVoxels();

ImagePlus rootImp = null;
int rootNjunctions = 0;
String warning = "";

// Create root image: a skeleton containing only branches inside a ROI
if (protectRoot) {
    rootImp = imp.duplicate();
    rootImp.setRoi(roi);
    for (int i=1; i<=rootImp.getNSlices(); i++) {
        ImageProcessor rootIp = rootImp.getStack().getProcessor(i);
        rootIp.setValue(0.0);
        rootIp.fillOutside(roi); //NB: This will fail if the ROI is the entire image
    }
    rootImp.killRoi();

    // Analyze root. We'll use n. of junctions as an identifier
    AnalyzeSkeleton_ root = new AnalyzeSkeleton_();
    root.setup("", rootImp);
    SkeletonResult rootResult = root.run(pruneChoice, false, false, origImp, true, false);
    if (rootResult.getNumOfTrees()>0) {
        int[] rootJunctions = rootResult.getJunctions();
        for (int i=0; i<rootJunctions.length; i++)
            rootNjunctions += rootJunctions[i];
    } else
        warning = "ROI contains no branches. Iteration was applied to entire image.\n";
}

if (protectRoot && rootNjunctions>0) {
    warning = "ROI contains "+ rootNJunctions +" branch point(s). Strahler counts will likely be inaccurate.\n";
    if (!IJ.showMessageWithCancel("Warning", warning +"Continue nevertheless?")) return;
}

// Use Z-projections to populate iteration stack when dealing with 3D skeletons
int nSlices = imp.getNSlices();
if (nSlices>1) {
    zp = new ZProjector(imp);
    zp.setMethod(ZProjector.MAX_METHOD);
    zp.setStartSlice(1);
    zp.setStopSlice(nSlices);
}

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

ImageStack newStack = new ImageStack( imp.getWidth(), imp.getHeight() );
int order = 1; int prevNjunctions = 0; boolean loop = true;
String msg = "\n*** Strahler analysis of "+ title +" ***\n"+ warning;

do {

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

    // Re-skeletonize image. This ensures 1) no disconnected pixels exist after pruning
    // and 2) that pruned structure is always a skeleton
    thin.run(ip);

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

    // We'll monitor the n. of remaining junctions
    SkeletonResult skelResult = skel.run(pruneChoice, false, false, origImp, true, false);
    int[] junctions = skelResult.getJunctions();
    if (junctions==null) {
        msg += "Iteration "+ order +": No branches were found! Revise parameters?\n";
        break;
    }
    int nJunctions = -rootNjunctions; // initialized to 0 if !protectRoot
    for (int i=0; i<junctions.length; i++)
        nJunctions += junctions[i];

    // Are there any loops in the skeleton?
    if (prevNjunctions==nJunctions) {
        msg += "Iteration "+ order +": Skeleton may contain loops at this step. Revise parameters?\n";
        break;
    }
    prevNjunctions = nJunctions;

    // Eliminate terminal branches
    skel.run(AnalyzeSkeleton_.NONE, true, false, null, true, false);

    if (protectRoot) { // Reinstate branches in root ROI
        (new ImageCalculator()).run("OR stack", imp, rootImp);
    }

    // Continue?
    loop = order<=maxPruning && nJunctions>0;

} while (loop);

// Create iteration stack. Add branch points to last slice
ImageProcessor pointsIp = ip.createProcessor(ip.getWidth(), ip.getHeight());
paintBranchPoints(pointsIp, junctionsList, 255);
newStack.addSlice("Branch points", pointsIp);
ImagePlus imp2 = new ImagePlus("StrahlerAnimation_"+ title, newStack);

// Create Strahler mask (16-bit image)
ZProjector 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 = ip3.convertToShortProcessor(false);
paintBranchPoints(ip3, junctionsList, 0); // Remove branch points
ImagePlus imp3 = new ImagePlus("StrahlerMask_"+ title, ip3);

// Retrieve counts using ParticleAnalyzer
ResultsTable tempRt = new ResultsTable();
ParticleAnalyzer pa = new ParticleAnalyzer(ParticleAnalyzer.SHOW_NONE, 0, tempRt, 0, Double.MAX_VALUE);
double currOrderCount = Double.NaN;
double nextOrderCount = Double.NaN;
for (int i=1; i<order; i++) {
    ip3.setThreshold(i, i, ImageProcessor.NO_LUT_UPDATE);
    pa.analyze(imp3);
    currOrderCount = (double)tempRt.getCounter();
    printResults(i, currOrderCount, currOrderCount/nextOrderCount);
    tempRt.reset();
    nextOrderCount = currOrderCount;
}

// Display images
Calibration cal = srcImp.getCalibration();
if (outChoice==OUT_IS || outChoice==OUT_ISCM) { // Display iteration stack
    imp2.setCalibration(cal);
    if (protectRoot) addRootToOverlay(imp2, roi, true);
    imp2.show();
}
if (outChoice==OUT_CM || outChoice==OUT_ISCM) { // Display color map
    ip3.resetThreshold();
    ip3.setMinAndMax(0, order-1);
    IJ.run(imp3, "Fire", "");
    IJ.run(imp3, "Calibration Bar...", " overlay"); // NB: Calibration bar clears pre-existing overlays. Bug?
    imp3.setCalibration(cal);
    if (protectRoot) addRootToOverlay(imp3, roi, false);
    imp3.show();
}

// Log messages
if (!msg.endsWith(" ***\n")) {
    if (protectRoot) msg += "Number of branch points (junctions) inside ROI: "+ rootNjunctions;
    IJ.log(msg);
}

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