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.
Originally inspired DAMAS by Thomas 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.

Started from version 8 of Iterative_Deconvolution

Version 1 6/2/04		"generalized inverse" for iteration, filter based
						on spectral content of psf

*/
/*	License:
	Copyright (c) 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_Deconvolution2 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", 4);
		double alpha = Prefs.get("iterativedevonvolution.alpha", 1000);
		boolean monotonic = Prefs.get("iterativedevonvolution.monotonic", true);
		boolean dB = Prefs.get("iterativedevonvolution.dB", false);

		GenericDialog gd = new GenericDialog("Iterative Deconvolution 2", 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("Alpha (suppress components < max/alpha in kernel)", alpha,0);
        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();
        alpha = 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.alpha", alpha);
		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;
		}
		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(dB)unDB(fhtA);
		if(operation == CONVOLVE){
			IJ.showStatus("Convolving: transform");
			fhtA.transform();
			FHT ipYCopy = new FHT(ipY);
			if(dB)unDB(ipYCopy);
			IJ.showStatus("Convolving: convolve in frequency domain");
			FHT ipX = convolve(ipYCopy,fhtA);
			IJ.showStatus("Convolving: creating new image");
			ipX.setRoi(bx,by,bw,bh);
			ImageProcessor ipXC = ipX.crop();
			if (dB) toDB(ipX);
			ipXC.setMinAndMax(0,0);
			ipXC.setColorModel(cmY);
			ImagePlus impX = new ImagePlus("Convolved",ipXC);
			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);
			//fhtY.swapQuadrants();
			float[] Y = (float[])fhtY.getPixels();
			if(dB)unDB(Y);
			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);
			ipX.setColorModel(cmY);
			ImagePlus impX = new ImagePlus("Deconvolved"+nIter+"_"+IJ.d2s(alpha),ipX);
			impX.show();
			//IJ.run("Fire");
			FHT 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 = new FHT(ipX);
				impX.setRoi(bx,by,bw,bh);  //set ROI just as a visual effect

				//Put Ax in XCopy
				ipXCopy = convolve(ipXCopy,fhtA);
				XCopy = (float[])ipXCopy.getPixels();

				//Put y - Ax in XCopy
				for (int ind = 0; ind < n; ind++){
					XCopy[ind] = (Y[ind] - XCopy[ind]/aSum);
				}

				//Find A-inverse times (y - Ax), and store in XCopy
				ipXCopy.transform();
				generalizedDivide(ipXCopy,fhtA,1/alpha,iN);
				ipXCopy.inverseTransform();
				ipXCopy.swapQuadrants();
				ipXCopy.multiply(aSum);

				//Increment X by that
				XCopy = (float[])ipXCopy.getPixels();

				for (int ind = 0; ind < n; ind++){
					X[ind] += XCopy[ind];
				}

				ipX.transform();
				filter(ipX,fhtA,1/(alpha*alpha),iN);
				ipX.inverseTransform();
				//ipX.swapQuadrants();

				//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);
			ipX.setMinAndMax(0,0);
			impX.setProcessor(null,ipX);
			impX.updateAndDraw();
			IJ.showStatus(nIter+" iterations complete.");
		}//if deconvolve
	}//run
	void unDB(ImageProcessor ip){
		float[] x = (float[])ip.getPixels();
		unDB(x);
	}
	void unDB(float[] x){
		double SCALE = 10/Math.log(10);
		int n = x.length;
		for (int i = 0; i < n; i++){
			x[i] = (float)Math.exp(x[i]/SCALE);
		}
	}
	void toDB(ImageProcessor ip){
		float[] x = (float[])ip.getPixels();
		toDB(x);
	}
	void toDB(float[] x){
		double SCALE = 10/Math.log(10);
		int n = x.length;
		for (int i = 0; i < n; i++){
			if(x[i] > 0.000000001)
				x[i] = (float)(SCALE*Math.log(x[i]));
			else
				x[i] = -90;
		}
	}
	FHT convolve(FHT ip, FHT fhtA){
		//FHT fhtIP = new FHT(ip);
		ip.transform();
		FHT result = ip.multiply(fhtA);
		result.inverseTransform();
		result.swapQuadrants();
		ip.inverseTransform();
		//ip.swapQuadrants();
		return result;
	}


    /** 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
     /** Returns the image resulting from the point by point Hartley division
        of this image by the specified image. Both images are assumed to be in
        the frequency domain. Division in the frequency domain is equivalent
        to deconvolution in the space domain.

        So far, this is Wanye Rasband's FHD divide.  The generalized part is that
        zero is returned for frequencies where the magnitude of the spectral
        component of h2 is less than alpha times the largest magnitude.  This is
        analogous to application of the generalized inverse in linear algebra.*/
    public void generalizedDivide(FHT arg, FHT fht, double alpha, int maxN) {
        int rowMod, cMod, colMod;
        double mag, h2e, h2o;
        float[] h1 = (float[])arg.getPixels();
        float[] h2 = (float[])fht.getPixels();
        float[] out = new float[maxN*maxN];
        double maxMag = 0;
		for (int r=0; r<maxN; r++) {
            rowMod = (maxN - r) % maxN;
            for (int c=0; c<maxN; c++) {
                colMod = (maxN - c) % maxN;
                mag =h2[r*maxN+c] * h2[r*maxN+c] + h2[rowMod*maxN+colMod] * h2[rowMod*maxN+colMod];
                if(mag > maxMag) maxMag = mag;
             }
        }
        double limit = alpha*maxMag;
		for (int r=0; r<maxN; r++) {
            rowMod = (maxN - r) % maxN;
            for (int c=0; c<maxN; c++) {
                colMod = (maxN - c) % maxN;
                mag =h2[r*maxN+c] * h2[r*maxN+c] + h2[rowMod*maxN+colMod] * h2[rowMod*maxN+colMod];

   				h2e = (h2[r*maxN+c] + h2[rowMod*maxN+colMod]);
                h2o = (h2[r*maxN+c] - h2[rowMod*maxN+colMod]);
                double tmp = (h1[r*maxN+c] * h2e - h1[rowMod*maxN+colMod] * h2o);
                out[r*maxN+c] = (float)(tmp/mag);

   				if(limit > 20.7*mag)
    				mag = 1000000000.;
    			else
   					mag *= (float)Math.exp(limit/mag);
               	out[r*maxN+c] = (float)(tmp/mag);

                /*if (mag<limit)
                    out[r*maxN+c] = 0;
                else{
					h2e = (h2[r*maxN+c] + h2[rowMod*maxN+colMod]);
                	h2o = (h2[r*maxN+c] - h2[rowMod*maxN+colMod]);
                	double tmp = (h1[r*maxN+c] * h2e - h1[rowMod*maxN+colMod] * h2o);
                	out[r*maxN+c] = (float)(tmp/mag);
				}*/
            }
        }
        arg.setPixels(out);
     }//generalizedDivide
     /** filter out spectral components of arg for which the
     corresponding component of fht is less than alpha times the
     max component of fht.*/
    public void filter(FHT arg, FHT fht, double alpha, int maxN) {
        int rowMod, cMod, colMod;
        double mag, h2e, h2o;
        float[] h1 = (float[])arg.getPixels();
        float[] h2 = (float[])fht.getPixels();
        float[] out = new float[maxN*maxN];
        double maxMag = 0;
		for (int r=0; r<maxN; r++) {
            rowMod = (maxN - r) % maxN;
            for (int c=0; c<maxN; c++) {
                colMod = (maxN - c) % maxN;
                mag =h2[r*maxN+c] * h2[r*maxN+c] + h2[rowMod*maxN+colMod] * h2[rowMod*maxN+colMod];
                if(mag > maxMag) maxMag = mag;
             }
        }
        double limit = alpha*maxMag;
		for (int r=0; r<maxN; r++) {
            rowMod = (maxN - r) % maxN;
            for (int c=0; c<maxN; c++) {
                colMod = (maxN - c) % maxN;
                mag =h2[r*maxN+c] * h2[r*maxN+c] + h2[rowMod*maxN+colMod] * h2[rowMod*maxN+colMod];
               // if (mag<limit)
                  //  h1[r*maxN+c] = 0;
				if(limit > 30*mag)
 					h1[r*maxN+c] = 0;
 				else
					h1[r*maxN+c] *= (float)Math.exp(-limit/mag);
			}
		}
    }//filter
	static public void showAbout(){
        IJ.showMessage( "About Iterative Deconvolution 2 ...",
                        "Iterative convolution and positive deconvolution\n" +
                        "Version 1, June 2, 2004\n" +
                      	"Author: Robert P. Dougherty, OptiNav, Inc.\n" +
                        "See http://www.optinav.com/ImageJplugins/Iterative-Deconvolution2.htm.");
    }
}//Iterative_Deconvolution