/* This file is based on the routine "seeded region growth" by Jarek Sacha
 *
 * Image/J Plugins package net.sf.ij_plugins.im3d.grow;
 * Copyright (C) 2002-2009 Jarek Sacha
 * Latest release available at http://sourceforge.net/projects/ij-plugins/
 *
 * 
 * The present file: SeededRegionGrowingForStacksWithMask.java is more or less a 
 * kind of GUI for the SRG plugin of Jarek Sacha
 *  
 * The idea of this file is to work through a stack slice by slice taking the seed points,
 * which are necessary for the SRG.java
 * from an other stack, which I call "label mask".
 *
 * The "label mask" can be created nearly automatically from grey stacks (for example
 * either by very conservative (= narrow) thresholding alone or by a combination of thresholding
 * and skeletonization...).
 *
 * To give you an idea:
 *
 * I use microcomputed data of teeth. Teeth are made of dentin and enamel which have different grey values.
 * I segment with a very conservative (narrow, not containing all greylevels) threshold to
 * make separate stacks for background, dentin and enamel.
 * Then I combine the masks according to my needs.
 * The thresholded images contain usually 0 for background and 255 for foreground.
 * I divide the thresholded images by 255 -> now the masked pixels have a value of 1
 * then I multiply the stack with 2 (for dentin) and 3 (for enamel).
 * The background, dentin and enamel stacks are combined into a new stack (= resulting "mask") with
 * ImageCalculator and "xor" (it works, did not think deeper until now, whether the approach is valid)
 *
 * According to my "convention" you should always use the value "1" as the value for the background label
 * as I will assign the value 1 to all slices which have less then 2 labels.
 *
 * The new stack is called "label mask" and the seedpoints are generated from this stack.
 *
 * The number of seedpoints will be adjusted to the minimum number of seeds within all labels.
 * This means, if we have 3 masks with 100, 1000 and 10000 points per label, then 100
 * will be the number of labels, which will be used for SRG.
 *
 * If one label has more than this minimum number of pixels, then the subset of this minimum number
 * will be selected equally distributed within this longer list.
 *
 * An new 8 bit stack results from the segmentation, which has the same labels as the label mask.
 *
 * Hint:
 *
 * I suggest to filter the raw data to improve the outcome, for example a median filter with a 3x3 or 5x5 kernel.
 *
 *
 * Personal disclaimer:
 *
 * I have no formal training in programming.
 * I am an autodidact with very limited time.
 * Usually I try to find solutions for similar
 * problems in the internet and then I adapt them
 * by trial and error to fit my purposes.
 *
 * I had to learn Java "on the job" (= dentist at the LMU, Munich, GER) because
 * I was tired to write grant applications, which take longer
 * then developing my own plugins for ImageJ.
 *
 * I would appreciate any help to improve my code!
 */

/**
 *
 * @author Karl-Heinz Kunzelmann
 * @author karl-heinz@kunzelmann.de
 * @author www.kunzelmann.de
 *
 * history:
 *
 * 26.09.2009 preliminary end of the project
 * 21.09.2009 started the project
 *
 * TODO: add a main method very similar to testBlobs() which allows commandline operation taking the filenames as args
 */

import ij.plugin.*;
import ij.*;
import ij.process.*;
import ij.ImagePlus;
import ij.process.ByteProcessor;
import ij.gui.*;
import java.io.IOException;
import java.util.logging.Level;
import java.util.logging.Logger;
import net.sf.ij_plugins.io.IOUtils;

import net.sf.ij_plugins.im3d.grow.*;

import java.awt.Point;
import java.util.Arrays;
import java.util.Set;
import java.util.TreeSet;
import net.sf.ij_plugins.grow.PointYXComparator;


public class SeededRegionGrowingForStacksWithMask implements PlugIn {


    // TODO long term plan - thread safe programming

    private static String title1 = "";
    private static String title2 = "";
    int maxNumberOfMaskPixels = 0;
    int speedUpFactor = 10;             // only every nth value of the label mask stack is processed
                                        // so far I did not really check the influence on computing time
                                        // or quality of the outcome

    int[] labels;                       // 0 is background of the label mask stack,
                                        // not included into the label array.
    
    int N;                              // used for array doubling
    Point[][] seeds;                    // array which is passed to the SRG routine of Jarek Sacha
    int maxSeedPoints = 0;              // will be automatically determined later
    int maxNumberOfLabelsInStack = 0;   // will be used to adjust brightness of result
    int selectMethodToGetSeedPoints = 1; // different methods to get seed points are available

    boolean preliminaryExit = false;        // if no stacks are open this flag is used for a nice exit
    boolean correctLabelMapFlag = false;    // not really necessary any longer, was used for developing
    boolean lessThenTwoLabels = true;       // if we have less then two labels, this flag is set true, else false

    ImagePlus img1;                     // grey value stack
    ImagePlus img2;                     // 8bit stack which serves as "label mask"
    ImagePlus img3;                     // result of SRG
    ImageStack stack1;
    ImageStack stack2;
    ImageStack stack3;

    ByteProcessor tempSlice;

    boolean debug = true;               // more or less meaningful infos for debugging
    boolean noDialog = false;           // used during development
                                        // if true, then no dialog is shown and the
                                        // files included to the source code are opened.

public void run(String arg) {

    coordinateAllTheWork();             // not really necessary....

}

public void coordinateAllTheWork() {
        
        if (noDialog){
            try {

                // I got tired to open my test images all the time via the message dialog

                img1 = IOUtils.openImage("C:/temp/Ubuntu-Transfer/00-grey.tif");
                img2 = IOUtils.openImage("C:/temp/Ubuntu-Transfer/00-mask.tif");
               
            } catch (IOException ex) {
                Logger.getLogger(SeededRegionGrowingForStacksWithMask.class.getName()).log(Level.SEVERE, null, ex);
            }
        }
        else {
            askForFilesDialog();
            if (preliminaryExit){
                return;
            }
        }

        img1.show();
        img2.show();
        stack1 = img1.getStack();
        stack2 = img2.getStack();

        // check whether grey stack and mask size have identical dimensions

        if (checkStackCompatibility(img1,img2)) {
            System.out.println("Stack size fits!");
        }
        else {
            System.out.println("Stack size does not fit!");
            return;
        }

        // new image stack for the results

        int widthStack = img1.getWidth();
        int heightStack = img1.getHeight();
        int depthStack = img1.getStackSize();

        stack3 = new ImageStack(widthStack,heightStack);
       
        // fill the stack with the segmentation result

        for (int i = 0; i < depthStack; i++){

            IJ.showProgress((double)(i+1)/(depthStack));

            //do something
            ImageProcessor greySlice = stack1.getProcessor(i+1); // Stacks start with 1
            ImageProcessor maskSlice = stack2.getProcessor(i+1);

            checkNumberOfMaskLabels(maskSlice);

            if (debug) System.out.println("fter checkNumberOfMaskLabels - slice: " + i);
            if (debug) System.out.println("fter checkNumberOfMaskLabels - labels.length/slice: "+ labels.length + "/" + i);
            if (debug) System.out.println("after checkNumberOfMaskLabels - lessThenTwoLabels: " + lessThenTwoLabels);

            // Check: are there more labels than background labels?
            
            if (!lessThenTwoLabels){
                try {

                    seeds = new Point[labels.length][maxSeedPoints];
                    if (debug) System.out.println("maxSeedP: " + maxSeedPoints);

                    stack3.addSlice(null,doRegionGrowing(greySlice, maskSlice));
                    

                } catch (Exception ex) {
                    System.out.println("Caught Exception in stack3.addSlice...");
                    Logger.getLogger(SeededRegionGrowingForStacksWithMask.class.getName()).log(Level.SEVERE, null, ex);
                }
            }
            else {

                // is it possible that some slices contain less then two labels
                // they are set to the background label then (= 1)

                ByteProcessor emptySlice = new ByteProcessor (widthStack, heightStack);
                emptySlice.add(1);  // 1 is the background label
                
                stack3.addSlice(null,emptySlice);
                

            }
        }

        img3 = new ImagePlus("Segmented_label_mask"+img1.getShortTitle(), stack3);
        img3.setDisplayRange(0, maxNumberOfLabelsInStack+1);         // even masks with just a few labels will be visible
        if (debug) System.out.println("maxNumberOfLabelsInStack: "+ maxNumberOfLabelsInStack);
        img3.show();
        return;
        
	}

private void askForFilesDialog(){

        //KHK: I took most of the next 30 lines from the ImageJ sources -> IJ.ImageCalculator.java

        int[] wList = WindowManager.getIDList();
        if (wList==null) {
                IJ.noImage();
                preliminaryExit = true;
                return;
        }

        String[] titles = new String[wList.length];

        for (int i=0; i<wList.length; i++) {
            ImagePlus imp = WindowManager.getImage(wList[i]);
            if (imp!=null)
                    titles[i] = imp.getTitle();
            else {
                titles[i]="";
            }
		}

        GenericDialog gd = new GenericDialog("SRG for Stacks", IJ.getInstance());
		String defaultItem;
        gd.addMessage("Please select two corresponding stacks:\n");
		if (title1.equals(""))
			defaultItem = titles[0];
		else
			defaultItem = title1;
		gd.addChoice("Grey level stack:", titles, defaultItem);

		if (title2.equals(""))
			defaultItem = titles[0];
		else
			defaultItem = title2;

        gd.addChoice("Label map stack (mask):", titles, defaultItem);

        String[] sedPoRedFactor = {"1", "10", "20", "50", "100"};
        defaultItem = sedPoRedFactor[1];

        gd.addChoice("SeedPoint reduction factor: ", sedPoRedFactor, defaultItem);

       	gd.showDialog();
		if (gd.wasCanceled())
			return;
		int index1 = gd.getNextChoiceIndex();
		title1 = titles[index1];
		int index2 = gd.getNextChoiceIndex();
		title2 = titles[index2];
        int index3 = gd.getNextChoiceIndex();
        speedUpFactor = Integer.parseInt(sedPoRedFactor[index3]);

        img1 = WindowManager.getImage(wList[index1]);
		img2 = WindowManager.getImage(wList[index2]);

        return;

}

private void checkNumberOfMaskLabels(ImageProcessor maskSlice){
    N = 0;

    labels = new int[1];
    
    int w = maskSlice.getWidth();
    int h = maskSlice.getHeight();


    maxNumberOfMaskPixels = w * h; // max. number of potential mask pixels = the whole image
       
        int[] pixelCounterPerLabel = new int[253];

        for (int y=0;y<h;y=y+speedUpFactor){
            for(int x=0;x<w;x=x+speedUpFactor){
                int temp=maskSlice.getPixel(x,y);
                // count labels:
                pixelCounterPerLabel[temp]=pixelCounterPerLabel[temp] + 1;


                // range of ByteProcessor is from 0 - 255, mask should be ByteProcessor
                // in SRG.java:  MAX_REGION_NUMBER = 253
                if (temp > 0 && temp < 253){
                    // if temp has a new label push it on the label stack
                    pushNewLabel(temp);
                }
            }
        }

        // determine the smallest number of defined label pixels per each label.
        // this number will determine the size of the seeds array, because it can be completely filled.
    

        maxSeedPoints = w*h;
        for (int c = 1; c < 253; c++)           // 0 is background and will not be evaluated, therefore c = 1
        {
            if (pixelCounterPerLabel[c] != 0){
                if (pixelCounterPerLabel[c]<maxSeedPoints){
                    maxSeedPoints = pixelCounterPerLabel[c];
                }
            }
        }

        // determine the maximum number of possible labels to
        // adjust the brightness of the resulting stack to the range of the available labels

        if (debug) System.out.println("in checkNumberOfMaskLabels: maxSeedPoints: "+maxSeedPoints);


        if (maxNumberOfLabelsInStack < labels.length){
            maxNumberOfLabelsInStack = labels.length;
        }

        if (debug){
            for (int i = 0; i < labels.length; i++){
                System.out.println("label list: "+ i + ": " + labels[i]);
            }
        }

 // The input mask can have any label number between 1 and 252
 // like for example: 1, 2, 4, 10, 17 etc.
 // The regionMaps after SRG is numbered continously starting with 1...
 // I process one slice after the other, but independently from each other
 // To ensure the same label for the same structure the label map has to be adjusted after SRG
 // to fit the original labels
 
 correctLabelMapFlag = true;   // no longer necessary, I wanted to save time in an earlier version,
                               // but it is more reliable to do correct every slice

 if (debug) System.out.println("boolean correctLabelMap: "+correctLabelMapFlag);

 return;
}


private void resizeLabelArray(int max){

    // this approach is taken from Sedgewick et al.: Algorithms and Data Structures pp. 560
    // the array is always just exactly the size we need

    // move stack to a new array of size max.
    int[] temp = new int[max];
    for (int i = 0; i < N; i++){
        temp[i] = labels[i];
    }
    labels = temp;
    return;
}

private void pushNewLabel(int newLabel){

    boolean addNewLabel = true;

    // this approach is taken from Sedgewick et al.: Algorithms and Data Structures pp. 560
    // the array is always just exactly the size we need

    // add new label to top of stack
    if (N == 0) {
            labels[N] = newLabel;
            N++;
            return;
    }


    for (int n = 0; n < N; n++){
        if (newLabel == labels[n]){
            addNewLabel = false;
        }
    }

    if (addNewLabel) {

             resizeLabelArray(labels.length+1);
             labels[N] = newLabel;
             Arrays.sort(labels);
             lessThenTwoLabels = false;
             N++;
    }

    return;

}

private boolean checkStackCompatibility(ImagePlus img1, ImagePlus img2){
    
        boolean sameSize;
        int w1 = img1.getWidth();
        int h1 = img1.getHeight();
        int w2 = img2.getWidth();
        int h2 = img2.getHeight();
        int size1 = img1.getStackSize();
        int size2 = img2.getStackSize();

	if ((size1>1 && size2>1 && size1!=size2) ||
            (w1 > 1 && w2 > 1 && w1 != w2)       ||
            (h1 > 1 && h2 > 1 && h1 != h2)){
                IJ.error("SeededRegionGrowingForStacks: ", "'Image1' and 'image2' must be stacks with the same\ndimensions");
                sameSize = false;
                return sameSize;
        }
        else {
            /*
             // TODO: I would like to check for ByteProcessor, dont know how
             // something like
             if (!(img2.getProcessor()).isBinary()){  // isBinary is wrong
               IJ.error("SeededRegionGrowingForStacks: ", "'Image2 (mask)' must be a 8bit stack");
               sameSize = false;
            }
            else*/
                sameSize = true;
            //}
        }

       return sameSize;

}


 public ByteProcessor doRegionGrowing(ImageProcessor greySlice, ImageProcessor maskSlice) throws Exception {

     // Load test image
        final ByteProcessor image = (ByteProcessor) greySlice;

        // TODO check whether SRG is possible at all
       
        // Setup region growing
        final SRG srg = new SRG();
        srg.setImage(image);

        if (selectMethodToGetSeedPoints == 0) {
            // select seed points in a linear sequence starting at 0,0 as soon as the given label is found
            extractSeedsArray(maskSlice);
        }
        else if (selectMethodToGetSeedPoints == 1){
            // select seed points as a equally distributed subset off all available seed points
            extractSeedsArrayUsingRandomSubset(maskSlice);
        }
        else if (selectMethodToGetSeedPoints == -1){
             // fix points for debugging
             setSampleSeedPointsForDebugging(maskSlice);
        }

        srg.setSeeds(seeds);
        srg.setNumberOfAnimationFrames(50);
        // Run growing
        srg.run();

        final ByteProcessor regionMask = srg.getRegionMarkers();

        if (debug){
            final ImageStack aniStack = srg.getAnimationStack();
            ImagePlus img4 = new ImagePlus("AniStack"+img1.getShortTitle(), aniStack);
            img4.setDisplayRange(0, maxNumberOfLabelsInStack+1);
            img4.show();
        }

        if (correctLabelMapFlag){
           correctLabelMap(regionMask);
        }
       return regionMask;
 }

 private ByteProcessor correctLabelMap(ByteProcessor regionMaskFromSRG){
        for (int h = 0; h < regionMaskFromSRG.getHeight(); h++){
            for (int w = 0; w < regionMaskFromSRG.getWidth(); w++){
                int temp = regionMaskFromSRG.get(w, h);
                if (temp != 0){
                   regionMaskFromSRG.set(w, h, labels[temp-1]);
                }
            }
        }
  return regionMaskFromSRG;
 }

 private void extractSeedsArray(ImageProcessor maskSlice){

     // The array size is already determined
     // each array element is individually filled with a Point(x,y)

     // only the first seed points up to maxSeedPoints are selected
     // isolated points may be ignored if they are located more to the lower right corner

     // loop for each label

     for (int u = 0; u < labels.length; u++){
            // height
            for (int y = 0; y < maskSlice.getHeight(); y=y+speedUpFactor) {
                // width
                for (int x = 0; x < maskSlice.getWidth(); x=x+speedUpFactor) {
                    int pixelValue = maskSlice.getPixel(x,y);             //for masks pixel value = label of mask
                    if (pixelValue == labels[u]){
                       
                        for (int v = 0; v < maxSeedPoints; v++){
                            seeds[u][v]=new Point(x,y);
                        }
                    }
                }
             }
     }
     return;
 }


 private void extractSeedsArrayUsingRandomSubset(ImageProcessor maskSlice){

     // The array size is already determined
     // each array element is individually filled with a Point(x,y)

     // loop for each label

     for (int u = 0; u < labels.length; u++){

         if (debug) System.out.println("label u: "+u
                + ", labels.length: " + labels.length
                + ", lessThenTwoLabelsFlag: " + lessThenTwoLabels);

       
        Point[] labelPointList = new Point[extractPointsPerLabel(u,maskSlice).length];
        labelPointList = extractPointsPerLabel(u,maskSlice);

        // A subset of equally distributed seed points is selected

        if (debug) System.out.println("lengthOfLabelPointList/maxSeedPoints: "+labelPointList.length+"/"+maxSeedPoints);

        for (int v = 0; v < maxSeedPoints ; v++ ){
            
            seeds[u][v]=new Point(labelPointList[v*(labelPointList.length/maxSeedPoints)].x,labelPointList[v*(labelPointList.length/maxSeedPoints)].y);
        }
     
     }
  return;
 }

private Point[][] setSampleSeedPointsForDebugging(ImageProcessor maskSlice){


   // KH: used just for development
   seeds = new Point[2][2];
   seeds[0][0]=new Point(1, 1);
   seeds[1][0]=new Point(60, 60);
   seeds[1][1]=new Point(65,65);

   System.out.println("seeds0: "+ seeds[0][0].getX()+", "+seeds[0][0].getY());
   System.out.println("seeds1: "+ seeds[1][0].getX()+", "+seeds[1][0].getY());
 
   return seeds;
}


 private Point[] extractPointsPerLabel(int label, ImageProcessor slice) {
        // if (debug) System.out.println("labelValue: "+ label);
        final Set<Point> points = new TreeSet<Point>(new PointYXComparator());

           for (int y = 0; y < slice.getHeight(); y=y+speedUpFactor) {
                for (int x = 0; x < slice.getWidth(); x=x+speedUpFactor) {

                    int pixelValue = slice.getPixel(x,y); //for masks pixel value = label of mask
                    if (pixelValue == labels[label]){ // background = 0 is not evaluated
                        //do something

                        points.add(new Point(x, y));

                        }
                    
                    }
                 }



        return points.toArray(new Point[points.size()]);
        }

}