//
// Frame_Selection.java
// Creator  : Nilanjan Ray (nray1@ualberta.ca)
// Date     : May 23, 2014
//
// Works on RGB or 8bit grayscale stacks or hyper stacks
// Works with single channel stacks or hyper stacks
// User input (1): which channels to use? "all" means R, G, and B colors are used
// User input (2): Number of images per time point (nImage). Default value is 3
// User input (3): Sigma for Gaussian smoothing. Default value 2. If input is less than 0.5, no smoothing is applied
//
// For a stack it selects one of nImage images. For a hyperstack it selects one of nImage images for each Zstack
// It modifies the input stack or the hyperstack by deleting images, which were not selected
// 
// Algorithm: Dynamic programming to enforce similarity between neighboring images
// Publication: 
//
import ij.*;
import ij.process.*;
import ij.gui.*;
import java.awt.*;
import ij.plugin.*;
import ij.plugin.frame.*;

public class Frame_Selection implements PlugIn {

    int nImage=3;
	int width=0;
	int height=0;
	int nSlice=0;
	int nFrame=0;
	int nTP=0;
	float sigmaGauss = 2;
	int[][] allIndices;
	
	public void run(String arg) {
		// get image from current window
		ImagePlus imp=WindowManager.getCurrentImage();
		if (null == imp){
			IJ.noImage();
 			return;
		} 

		// Dialogbox to get some user inputs
		GenericDialog gd = new GenericDialog("Frame Selection");
		final String[] imChannels = new String[] {"red", "green", "blue","all"} ;
		gd.addChoice("Channels to use: ", imChannels, imChannels[3]);
		gd.addNumericField("# of images / timepoint: ", nImage, 0);
		gd.addNumericField("Sigma for Gaussian smooothong: ", sigmaGauss, 1);
		gd.showDialog();
		if (gd.wasCanceled()) return;
		final int itype = gd.getNextChoiceIndex();  
		nImage = (int)gd.getNextNumber();
		sigmaGauss = (float)gd.getNextNumber();
		if(nImage<2){
			IJ.showMessage("Frame_Selection","# images / timepoint should be >= 2");
			return;
		}
		
		// get dimensions of the stack / hyperstack and set variables
		int[] dim = imp.getDimensions();		
		if(dim[2] != 1){
			IJ.showMessage("Frame_Selection","Requires # channels to be 1");
			return;
		}
		width=dim[0];
		height=dim[1];
		nSlice=dim[3];
		nFrame=dim[4];
		nTP = nSlice / nImage;
		
		// get stack
		ImageStack stack=imp.getStack();
		if(stack.getBitDepth() != 24 && stack.getBitDepth() != 8){ // suppots RGB and 8-bit grayscale image
			IJ.showMessage("Frame_Selection","Unsupported image format");
			return;			
		}
		
		// preprocess stack: most expensive component when blurring is applied
		ImageStack newStack = preprocess(stack);
		
		// indices for image selection
		allIndices = new int[nTP][nFrame];
		
		ImageStack[] stackChannels = null;
		float[][][] cost = new float[nImage][nImage][nTP-1];
		if(stack.getBitDepth()==24) // RGB image
			stackChannels = ChannelSplitter.splitRGB(newStack,false);
		for(int nf=0; nf<nFrame; nf++){
			// build cost function
			if(stack.getBitDepth()==24){ // RGB image
				if(itype<3){
					cost = buildCost(stackChannels[itype],nf);
				}else{
					float[][][] costr = buildCost(stackChannels[0],nf);
					float[][][] costg = buildCost(stackChannels[1],nf);
					float[][][] costb = buildCost(stackChannels[2],nf);
					for(int m=0;m<nImage;m++)
						for(int n=0;n<nImage;n++)
							for(int i=0;i<nTP-1;i++)
								cost[m][n][i] = (costr[m][n][i] + costg[m][n][i] + costb[m][n][i])/3;
				}
			}else{
				if(stack.getBitDepth()==8){ // 8 bit image
					cost = buildCost(newStack,nf);
				}
			}
			// call DP optimizer to get indices of good images
			int[] indices = dynamicProg(cost);
			for(int n=0; n<nTP; n++)
				allIndices[n][nf] = indices[n];
			IJ.showProgress(nf+1, nFrame);
		}
		
		// delete not-so-good images from the stack
		int deleteCount = 0;
		for(int nf=0; nf<nFrame; nf++){
			for(int n=0;n<nTP;n++){
				for(int i=0;i<nImage;i++){
					if(i != allIndices[n][nf]){
						stack.deleteSlice(nf*nTP*nImage + n*nImage + i+1-deleteCount);
						deleteCount += 1;
					}
				}
			}
		}
		
		// update display
		imp.setDimensions(1,nTP,nFrame);
		imp.setOpenAsHyperStack(true);
		imp.show(); 
		IJ.showMessage("Frame_Selection","Images are selected!");
	}
	private ImageStack preprocess(ImageStack stack){
		// preprocess stack: create a new stack of 1/resizeFactor the width and the height of the input stack
		final int resizeFactor = 4;
		ImageStack stackOut = new ImageStack((int)(stack.getWidth()/resizeFactor),(int)(stack.getHeight()/resizeFactor));
		for(int n=1; n<=nSlice*nFrame; n++){
			ImageProcessor im = stack.getProcessor(n);
			ImageProcessor imDuplicate = im.duplicate();
			if(sigmaGauss>=0.5)
				imDuplicate.blurGaussian(sigmaGauss); // Gaussian smoothing filter to suppress noise
			imDuplicate.setInterpolationMethod(ImageProcessor.BILINEAR);
			ImageProcessor im1 = imDuplicate.resize((int)(stack.getWidth()/resizeFactor),(int)(stack.getHeight()/resizeFactor));
			stackOut.addSlice(im1);
			IJ.showProgress(n, nSlice*nFrame);
		}
		return stackOut;
	}
	private float[][][] buildCost(ImageStack stack, int whichFrame){
				
		float[][][] cost = new float[nImage][nImage][nTP-1];
		for(int n=0; n<nTP-1; n++){
			for(int i=0;i<nImage;i++){
				ImageProcessor im1 = stack.getProcessor(whichFrame*nTP*nImage + n*nImage + i+1);
				ImageProcessor ip1 = null;
				if(whichFrame>0)
					ip1 = stack.getProcessor((whichFrame-1)*nTP*nImage + n*nImage + allIndices[n][whichFrame-1]+1);
				for(int j=0;j<nImage;j++){
					ImageProcessor im2 = stack.getProcessor(whichFrame*nTP*nImage + (n+1)*nImage + j+1);
					cost[i][j][n] = mad((byte[])im1.getPixels(), (byte[])im2.getPixels(), im1.getWidth()*im1.getHeight());
					if(whichFrame>0){
						ImageProcessor ip2 = stack.getProcessor((whichFrame-1)*nTP*nImage + (n+1)*nImage + allIndices[n+1][whichFrame-1]+1);
						float c1 = mad((byte[])im1.getPixels(), (byte[])ip1.getPixels(), im1.getWidth()*im1.getHeight());
						float c2 = mad((byte[])im2.getPixels(), (byte[])ip2.getPixels(), im1.getWidth()*im1.getHeight());
						cost[i][j][n] += (c1+c2)/2;
					}
				}
			}
		}
		return cost;
	}
	private float mad(byte[] im1, byte[] im2, int len){
		float pixelMad=0;
		for(int i=0; i<len; i++){
			int diff = (int)im1[i]-(int)im2[i];
			if(diff<0){ diff = -diff;}
			pixelMad += (float)diff/(float)len;
		}
		return pixelMad;
	}
	private int[] dynamicProg(float[][][] cost){
		int N = nTP;
		int n = nImage;

		float[][] vf = new float[n][N]; // value functions
		int[][] vi = new int[n][N]; // indices for minimum values
		int[] indices = new int[N]; // return indices
		
		for(int j=0; j<N-2;j++){
			for(int i=0;i<n;i++){
				float mincost = Float.MAX_VALUE;
				int minvi = 0;
				for(int k=0;k<n;k++){
					float thiscost = cost[k][i][j] + vf[k][j];
					if(thiscost<mincost){
						mincost = thiscost;
						minvi = k;
					}
				}
				vf[i][j+1] = mincost;
				vi[i][j+1] = minvi;
			}
		}
		float minval=Float.MAX_VALUE;
		int mink=0;
		for(int k=0;k<n;k++){
			float thisval = vf[k][N-1];
			if(thisval<minval){
				minval=thisval;
				mink=k;
			}
		}
		indices[N-1] = mink;
		for(int j=N-2;j>=0;j--){
			indices[j] = vi[indices[j+1]][j+1];
		}
		return indices;
	}
}