import ij.*;
import ij.gui.*;
import ij.plugin.*;
import ij.process.*;
import ij.plugin.filter.Convolver;

/*Image convolution and postive deconvolution.  Bob Dougherty 5/18/04.
Loosely based on DAMAS by Tomas F. Brooks and William M. Humphreys, Jr.
AIAA Paper 2004-2954

Some code is from Wayne Rasband's FFTFilter, which is, in turn, based on
Joachim Walter's FFT Filter plugin.

Version 0 5/18/04
Version 1 5/18/04 pm 	added low pass filter
Version 2 5/19/04 		generalized image type and dimensions
Version 3 5/26/04		fixed a bug that excluded 8-bit psf images
Version 4 5/26/04		saves user inputs in the preferences file
Version 5 5/28/04		"monotonic" option
Version 6 5/28/04		fixed a bug that produced rotations when the image is
						much larger than the kernel
Version 7 5/31/04		limits processing to the ROI of the image and the psf
Version 7.1 5/31/04		do not copy the image or psf unnecessarily if there is no ROI
Version 7.2 5/31/04		improved accuracy of centering of psf
Version 8 6/1/04		cleaned up code, reduced memory requirements
Version 9 2/0/05		repaired dB option

*/
/*	License:
	Copyright (c) 2004, 2005, OptiNav, Inc.
	All rights reserved.

	Redistribution and use in source and binary forms, with or without
	modification, are permitted provided that the following conditions
	are met:

		Redistributions of source code must retain the above copyright
	notice, this list of conditions and the following disclaimer.
		Redistributions in binary form must reproduce the above copyright
	notice, this list of conditions and the following disclaimer in the
	documentation and/or other materials provided with the distribution.
		Neither the name of OptiNav, Inc. nor the names of its contributors
	may be used to endorse or promote products derived from this software
	without specific prior written permission.

	THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
	"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
	LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
	A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
	CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
	EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
	PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
	PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
	LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
	NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
	SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
public class Iterative_Deconvolution implements PlugIn{
    private static final int CONVOLVE=0, DECONVOLVE=1;
   	private static String[] ops = {"Convolve", "Deconvolve"};
 	public void run(String arg){
        if (arg.equalsIgnoreCase("about")){
            showAbout();
            return;
        }
		if (IJ.versionLessThan("1.32c"))
			return;
 		int[] wList = WindowManager.getIDList();
        if (wList == null){
            IJ.noImage();
            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] = "";
        }
		String titleImage = Prefs.get("iterativedevonvolution.titleImage", titles[0]);
		int imageChoice = 0;
        for (int i = 0; i < wList.length; i++){
			if(titleImage.equals(titles[i])){
				imageChoice = i;
				break;
			}
		}
		String titlePSF = Prefs.get("iterativedevonvolution.titlePSF", titles[0]);
		int psfChoice = 0;
        for (int i = 0; i < wList.length; i++){
			if(titlePSF.equals(titles[i])){
				psfChoice = i;
				break;
			}
		}
		int operation = (int)Prefs.get("iterativedevonvolution.operation", DECONVOLVE);
		int nIter = (int)Prefs.get("iterativedevonvolution.nIter", 16);
		double filterSmallDia = Prefs.get("iterativedevonvolution.filterSmallDia", 2);
		boolean monotonic = Prefs.get("iterativedevonvolution.monotonic", true);
		boolean dB = Prefs.get("iterativedevonvolution.dB", false);

		GenericDialog gd = new GenericDialog("Iterative Deconvolution", IJ.getInstance());
		gd.addChoice("Operation", ops, ops[operation]);
        gd.addChoice("Image",titles,titles[imageChoice]);
        gd.addChoice("Point Spread Function (Kernel)",titles,titles[psfChoice]);
		gd.addNumericField("Number of iterations (applies to deconvolution)", nIter,0);
		gd.addNumericField("LP filter diameter, pixels (applies to deconvolution)", filterSmallDia,1);
        gd.addCheckbox("Monotonic deconvolution option",monotonic);
       	gd.addCheckbox("Image, psf and result in dB",dB);
        gd.showDialog();

        if (gd.wasCanceled())
            return;
		operation = gd.getNextChoiceIndex();
        ImagePlus impY = WindowManager.getImage(wList[gd.getNextChoiceIndex()]);
        ImagePlus impA = WindowManager.getImage(wList[gd.getNextChoiceIndex()]);
        nIter = (int)gd.getNextNumber();
        filterSmallDia = gd.getNextNumber();
        monotonic = gd.getNextBoolean();
       	dB = gd.getNextBoolean();

		Prefs.set("iterativedevonvolution.operation", operation);
		Prefs.set("iterativedevonvolution.titleImage", impY.getTitle());
		Prefs.set("iterativedevonvolution.titlePSF", impA.getTitle());
		Prefs.set("iterativedevonvolution.nIter", nIter);
		Prefs.set("iterativedevonvolution.filterSmallDia", filterSmallDia);
		Prefs.set("iterativedevonvolution.monotonic", monotonic);
		Prefs.set("iterativedevonvolution.dB", dB);

        ImageProcessor ipY = null;
        ImageProcessor ipA = null;
		if(impY.getRoi()!=null){
			//IJ.write("image has ROI");
			ipY = impY.getProcessor().crop();
		}else{
			//IJ.write("image does not have ROI");
			ipY = impY.getProcessor();
		}

		if(impA.getRoi()!=null){
			//IJ.write("kernel has ROI");
			ipA = impA.getProcessor().crop();
		}else{
			//IJ.write("kernel does not have ROI");
			ipA = impA.getProcessor();
		}

 		if(((ipY instanceof ColorProcessor)||(ipA instanceof ColorProcessor))){
			IJ.showMessage("RGB images are not currently supported.");
			return;
		}

		double minA = 0;
		double minY = 0;
		if(dB){
			minA = unDB(ipA);
			minY = unDB(ipY);
		}

		java.awt.image.ColorModel cmY = ipY.getColorModel();
		int kw = ipA.getWidth();
		int kh = ipA.getHeight();
		int bw = ipY.getWidth();
		int bh = ipY.getHeight();
		int maxN = kw;
		if(kh > maxN)maxN = kh;
		if(bw > maxN)maxN = bw;
		if(bh > maxN)maxN = bh;
		//Expand this to a power of 2 that is at least 1.5* as large, to avoid wrap effects
        int iN=2;
       	while(iN<1.5 * maxN) iN *= 2;
        // fit blurred image and kernel into power of two square images
		ipA = insert(ipA, iN, iN);
		int bx = (int) Math.round((iN - bw)/2.0);
		int by = (int) Math.round((iN - bh)/2.0);
		ipY = tileMirror(ipY, iN, iN, bx, by);


		int n = iN*iN;
		FHT fhtA = new FHT(ipA);
		if(operation == CONVOLVE){
			IJ.showStatus("Convolving: transform");
			fhtA.transform();
			ImageProcessor ipYCopy = ipY.duplicate();
			IJ.showStatus("Convolving: convolve in frequency domain");
			ImageProcessor ipX = convolve(ipYCopy,fhtA);
			IJ.showStatus("Convolving: creating new image");
			ipX.setRoi(bx,by,bw,bh);
			ipX = ipX.crop();
			if (dB) toDB(ipX,-90);
			ipX.setMinAndMax(0,0);
			ipX.setColorModel(cmY);
			ImagePlus impX = new ImagePlus("Convolved",ipX);
			impX.show();
			IJ.showStatus("done");
		}else{
			float[] A = (float[])fhtA.getPixels();
			float aSum = 0;
			for (int j = 0; j < iN; j++){
				for (int i = 0; i < iN; i++){
					aSum += A[i + iN*j];
				}
			}
			fhtA.transform();
			FHT fhtY = new FHT(ipY);
			float[] Y = (float[])fhtY.getPixels();
			float[] min = new float[n];
			float[] max = new float[n];
			for (int ind = 0; ind < n; ind++){
				min[ind] = 0;
				max[ind] = Float.MAX_VALUE;
			}
			if(monotonic){
				ImageProcessor ipAY = convolve(fhtY,fhtA);
				float[] AY = (float[])ipAY.getPixels();
				for (int ind = 0; ind < n; ind++){
					float convolvedOnce = Y[ind];
					float convolvedTwice = AY[ind]/aSum;
					if(convolvedTwice >= convolvedOnce)
						max[ind] = convolvedOnce;
					if(convolvedTwice <= convolvedOnce)
						min[ind] = convolvedOnce;
				}
			}
			//FHT ipX = new FHT(ipY);

			ImageProcessor ipXtemp = new FloatProcessor(iN,iN);
			FHT ipX = new FHT(ipXtemp);
			ipX.setColorModel(cmY);

			ImagePlus impX = new ImagePlus("Deconvolved"+nIter+"_"+IJ.d2s(filterSmallDia),ipX);
			impX.show();
			//IJ.run("Fire");
			ImageProcessor ipXCopy = null;
			float[] XCopy = null;
			float[] X = (float[])ipX.getPixels();;
			for (int iter = 0; iter < nIter; iter++){
				IJ.showProgress((float)iter/nIter);
				IJ.showStatus("Starting iteration "+(iter+1)+" of "+nIter);
				impX.killRoi();
				ipXCopy = ipX.duplicate();
				impX.setRoi(bx,by,bw,bh);  //set ROI just as a visual effect
				ipXCopy = convolve(ipXCopy,fhtA);
				XCopy = (float[])ipXCopy.getPixels();
				for (int ind = 0; ind < n; ind++){
					X[ind] += (Y[ind] - XCopy[ind]/aSum);
				}
				filterSmall(ipX,filterSmallDia);
				X = (float[])ipX.getPixels();
				for (int ind = 0; ind < n; ind++){
					if(X[ind] < min[ind])X[ind] = min[ind];
					if(X[ind] > max[ind])X[ind] = max[ind];
				}
				ipX.setMinAndMax(0,0);
				impX.updateAndDraw();
			}
			ipX.setRoi(bx,by,bw,bh);
			ipX = new FHT(ipX.crop());
			ipX.multiply(1/aSum);
			if (dB) toDB(ipX,-90);
			ipX.setMinAndMax(0,0);
			impX.setProcessor(null,ipX);
			impX.updateAndDraw();
			IJ.showStatus(nIter+" iterations complete.");
		}//if deconvolve
		if(dB){
			ipA = impA.getProcessor();
			ipY = impY.getProcessor();
			toDB(ipA,minA);
			toDB(ipY,minY);
		}
	}//run
	double unDB(ImageProcessor ip){
		float[] x = (float[])ip.getPixels();
		return unDB(x);
	}
	double unDB(float[] x){
		double SCALE = 10/Math.log(10);
		int n = x.length;
		double result = Float.MAX_VALUE;
		for (int i = 0; i < n; i++){
			if(x[i] < result) result = x[i];
			x[i] = (float)Math.exp(x[i]/SCALE);
		}
		return result;
	}
	void toDB(ImageProcessor ip, double minDB){
		float[] x = (float[])ip.getPixels();
		toDB(x, minDB);
	}
	void toDB(float[] x, double minDB){
		double SCALE = 10/Math.log(10);
		double minVal = Math.exp(minDB/SCALE);
		int n = x.length;
		for (int i = 0; i < n; i++){
			if(x[i] > minVal)
				x[i] = (float)(SCALE*Math.log(x[i]));
			else
				x[i] = (float)minDB;
		}
	}
	ImageProcessor convolve(ImageProcessor ip, FHT fhtA){
		FHT fhtIP = new FHT(ip);
		fhtIP.transform();
		FHT result = fhtIP.multiply(fhtA);
		result.inverseTransform();
		result.swapQuadrants();
		return (FloatProcessor)result;
	}

    /*
    filterSmall: up to which size are small structures suppressed?
    filterSmall is given as fraction of the image size
                in the original (untransformed) image.

    This is adapted from Wayne Rasband's FFTFilter, which is, in turn, based on
  	Joachim Walter's FFT Filter plugin at "http://rsb.info.nih.gov/ij/plugins/fft-filter.html".
	*/
    void filterSmall(FHT ip,double filterSmallDia) {

        int maxN = ip.getWidth();
        double filterSmall = filterSmallDia / maxN;
		ip.transform();

        float[] fht = (float[])ip.getPixels();
        float[] filter = new float[maxN*maxN];
        for (int i=0; i<maxN*maxN; i++)
            filter[i]=1f;

        int row;
        int backrow;
		float rowFactSmall;

        int col;
        int backcol;
        float factor;
        float colFactSmall;

        // calculate factor in exponent of Gaussian from filterSmall
        double scaleSmall = filterSmall*filterSmall;

        // loop over rows
        for (int j=1; j<maxN/2; j++) {
            row = j * maxN;
            backrow = (maxN-j)*maxN;
             rowFactSmall = (float) Math.exp(-(j*j) * scaleSmall);


            // loop over columns
            for (col=1; col<maxN/2; col++){
                backcol = maxN-col;
                 colFactSmall = (float) Math.exp(- (col*col) * scaleSmall);
                factor =  rowFactSmall*colFactSmall;

                fht[col+row] *= factor;
                fht[col+backrow] *= factor;
                fht[backcol+row] *= factor;
                fht[backcol+backrow] *= factor;
                filter[col+row] *= factor;
                filter[col+backrow] *= factor;
                filter[backcol+row] *= factor;
                filter[backcol+backrow] *= factor;
            }
        }

        //process meeting points (maxN/2,0) , (0,maxN/2), and (maxN/2,maxN/2)
        int rowmid = maxN * (maxN/2);
        rowFactSmall = (float) Math.exp(- (maxN/2)*(maxN/2) * scaleSmall);

        fht[maxN/2] *= rowFactSmall; // (maxN/2,0)
        fht[rowmid] *= rowFactSmall; // (0,maxN/2)
        fht[maxN/2 + rowmid] *= rowFactSmall*rowFactSmall; // (maxN/2,maxN/2)
        filter[maxN/2] *=rowFactSmall; // (maxN/2,0)
        filter[rowmid] *= rowFactSmall; // (0,maxN/2)
        filter[maxN/2 + rowmid] *= rowFactSmall*rowFactSmall; // (maxN/2,maxN/2)

        //loop along row 0 and maxN/2
		rowFactSmall = (float) Math.exp(- (maxN/2)*(maxN/2) * scaleSmall);
        for (col=1; col<maxN/2; col++){
            backcol = maxN-col;
			colFactSmall = (float) Math.exp(- (col*col) * scaleSmall);
			fht[col] *= colFactSmall;
			fht[backcol] *= colFactSmall;
			fht[col+rowmid] *= colFactSmall*rowFactSmall;
			fht[backcol+rowmid] *= colFactSmall*rowFactSmall;
			filter[col] *= colFactSmall;
			filter[backcol] *= colFactSmall;
			filter[col+rowmid] *= colFactSmall*rowFactSmall;
			filter[backcol+rowmid] *= colFactSmall*rowFactSmall;
        }

        // loop along column 0 and maxN/2
        colFactSmall = (float) Math.exp(- (maxN/2)*(maxN/2) * scaleSmall);
        for (int j=1; j<maxN/2; j++) {
            row = j * maxN;
            backrow = (maxN-j)*maxN;
			rowFactSmall = (float) Math.exp(- (j*j) * scaleSmall);

			fht[row] *= rowFactSmall;
			fht[backrow] *= rowFactSmall;
			fht[row+maxN/2] *= rowFactSmall*colFactSmall;
			fht[backrow+maxN/2] *= rowFactSmall*colFactSmall;
			filter[row] *= rowFactSmall;
			filter[backrow] *= rowFactSmall;
			filter[row+maxN/2] *= rowFactSmall*colFactSmall;
			filter[backrow+maxN/2] *= rowFactSmall*colFactSmall;
         }
 		ip.inverseTransform();
    }//filterSmall

    /** Puts imageprocessor (ROI) into a new imageprocessor of size width x height y at position (x,y).
    The image is mirrored around its edges to avoid wrap around effects of the FFT.
    This is borrowed from Wayne Rasband's FFTFilter, which is, in turn, based on
  	Joachim Walter's FFT Filter plugin at "http://rsb.info.nih.gov/ij/plugins/fft-filter.html".
    */
    ImageProcessor tileMirror(ImageProcessor ip, int width, int height, int x, int y) {

        if (x < 0 || x > (width -1) || y < 0 || y > (height -1)) {
            IJ.error("Image to be tiled is out of bounds.");
            return null;
        }
        ImageProcessor ipout = ip.createProcessor(width, height);

        ImageProcessor ip2 = ip.duplicate();
        int w2 = ip2.getWidth();
        int h2 = ip2.getHeight();

        //how many times does ip2 fit into ipout?
        int i1 = (int) Math.ceil(x / (double) w2);
        int i2 = (int) Math.ceil( (width - x) / (double) w2);
        int j1 = (int) Math.ceil(y / (double) h2);
        int j2 = (int) Math.ceil( (height - y) / (double) h2);

        //tile
        if ( (i1%2) > 0.5)
            ip2.flipHorizontal();
        if ( (j1%2) > 0.5)
            ip2.flipVertical();

        for (int i=-i1; i<i2; i += 2) {
            for (int j=-j1; j<j2; j += 2) {
                ipout.insert(ip2, x-i*w2, y-j*h2);
            }
        }

        ip2.flipHorizontal();
        for (int i=-i1+1; i<i2; i += 2) {
            for (int j=-j1; j<j2; j += 2) {
                ipout.insert(ip2, x-i*w2, y-j*h2);
            }
        }

        ip2.flipVertical();
        for (int i=-i1+1; i<i2; i += 2) {
            for (int j=-j1+1; j<j2; j += 2) {
                ipout.insert(ip2, x-i*w2, y-j*h2);
            }
        }

        ip2.flipHorizontal();
        for (int i=-i1; i<i2; i += 2) {
            for (int j=-j1+1; j<j2; j += 2) {
                ipout.insert(ip2, x-i*w2, y-j*h2);
            }
        }

        return ipout;
    }//tile mirror
    /** Puts imageprocessor (ROI) into a new imageprocessor of size width x height y at position (x,y).
    Zero pad outside.  Try to put CM of the kernel in the center of the image. */
    ImageProcessor insert(ImageProcessor ip, int width, int height) {
        ImageProcessor ipout = new FloatProcessor(width, height);
        float xCM = 0;
        float yCM = 0;
        float area = 0;
        int w = ip.getWidth();
        int h = ip.getHeight();
		ImageProcessor ipf = ip.convertToFloat();
        float[] pixels = (float[])ipf.getPixels();
        for (int j = 0; j < h; j++){
			for (int i = 0; i < w; i++){
				float data = pixels[i + w*j];
				xCM += i*data;
				yCM += j*data;
				area += data;
			}
		}
		xCM /= area;
		yCM /= area;
		int x = (int)(width/2. - xCM + 0.5);
		int y = (int)(height/2. - yCM + 0.5);

		if(x >= width) x = width-1;
		if(x < 0) x = 0;
		if(y >= height) y = height-1;
		if(y < 0) y = 0;
		ipout.insert(ipf, x, y);
        return ipout;
    }//insert
    static public void showAbout(){
        IJ.showMessage( "About Iterative Deconvolution ...",
                        "Iterative convolution and positive deconvolution\n" +
                        "Version 8, May 31, 2004\n" +
                      	"Author: Robert P. Dougherty, OptiNav, Inc.\n" +
                        "See http://www.optinav.com/ImageJplugins/Iterative-Deconvolution.htm.");
    }

}//Iterative_Deconvolution