package vnt;

/**
 * Lighting_Correction.java - v1.0
 * Began: June 27, 2005
 * Last Updated: July 28, 2005 
 * 
 * Copyright (C) 2005-2006 - Michael D. Miller - mdm162@truman.edu
 * Truman State University
 * 100 E. Normal
 * Kirksville, MO - 63501
 * 
 * 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; either version 2
 * of the License, or (at your option) any later version.
 * 
 * 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.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
 */

import ij.*;
import ij.gui.*;
import java.awt.*;
import java.awt.Color.*;
import ij.plugin.filter.PlugInFilter;
import ij.process.*;
import java.io.*;
import ij.io.*;
import java.lang.String.*;
import java.lang.Math.*;
import ij.plugin.filter.*;

import Jama.Matrix;

/**
 * <p>This plug-in is meant to be used for images with a no-reducing
 * lense. Any images with a single centered light should use this.</p>
 * 
 * <p>Accepts a grayscale image as input and fits an n-th order 
 * polynomial surface by treating color as a z value. The fit surface 
 * is used to adjust lighting irregularities accordingly.</p>
 * 
 * <p>This class relies on JAMA, a public domain Java Matrix package
 * for all least-squares regression matrix calculations, freely
 * available at http://math.nist.gov/javanumerics/jama/.</p>
 *  
 * @author Michael Miller - Truman State University
 * @version 1.0
 * @since 1.0
 */
public class Lighting_Correction extends VascularNetworkToolkit implements PlugInFilter {

    /**
     * The 'n' order of the polynomial surface. Also read as degree.
     */
    private int order;
    
    /**
     * The number of terms for the given polynomial surface.
     * orderTerms = ((order+1)*(order+1)+(order+1))/2;
     * 
     * Ex: If order = 3, then the surface is given to be
     * f(x,y) = a + b*x + c*y + d*x^2 + e*x*y + f*y^2 + g*x^3
     *             + h*x*y^2 + i*x^2*y + j*y^3
     * 
     * Here orderTerms is 10, or given by the equation
     * [(order+1)*(order+1)+(order+1)]/2 = 10
     */
    private int orderTerms;

    /**
     * The array that stores the terms for the function f(x,y).
     * Ex: If order = 3, then the surface is given to be
     * f(x,y) = a + b*x + c*y + d*x^2 + e*x*y + f*y^2 + g*x^3
     *             + h*x*y^2 + i*x^2*y + j*y^3
     * Here, coefficients[0] == a, coefficients[7] == h, etc.
     */
    private double [] coefficients;
    
    /**
    * Specifies the preconditions for the plug-in.
    * If this method succeeds then run() is called.
    *
    * <p>Pre: ImageJ is running and an 8-bit grayscale image is open. The plug-in was just activated.
    * <br />Post: Either an argument was processed, the image was not saved to a local folder, or the plug-in is cleared to run on the image.
    * @param arg Required by the interface. The argument list passed to the plug-in.
    * @param imp Required by the interface. The ImagePlus datastructure holding (among other things) information to grab path and filename.
    * @return If DONE is returned, ImageJ quits without run()'ing the plug-in. Otherwise, the plug-in signals to ImageJ that this plugin only handles 8-bit (256 grayscale) and will not change the original image.
    * @see #run
    */
    public int setup(String arg, ImagePlus ip) {
        if (arg.equals("about")) { 
            showAbout("Lighting Correction", "  * Least-squares fits a quadratic surface to the image to adjust for point lighting.  (Copyright 2005. Michael Miller mdm162@truman.edu)");
            return DONE; // exit without run()
        }

        // will only save if this succeeds
        getFileInformation(ip);
                    
        return DOES_8G+NO_CHANGES; // success, run()
    }

    /**
     * The executed method when the plug-in is successfully called.
     * 1) Performs a n'th order polynomial surface least-squares fit on the original image.
     * 2) Creates a new image and displays the polynomial surface fit.
     * 3) Copies the old image onto a second new image with light-corrected values.
     *
     * <p>Pre: The image was cleared to run by the setup() method.
     * <br />Post:  The image is processed by the lighting correction routines. A new corrected image is drawn. 
     * @param bp Required by the interface. The access information to the original image.
     * @see #setup
     * @see #setDegree
     * @see #calculateSurfaceCoefficients
     * @see #drawSurface
     * @see #drawCorrectedImage
     */
    public void run(ImageProcessor bp){         
        // get the width and height information
        getDimensions(bp);
        // set the maximum order in the polynomial surface
        setDegree(lightingCorrectionPolynomialOrder);// ie this is a variable-th degree polynomial surface
        // take the image, and find it's least squares regression coefficients for the specified polynomial
        calculateSurfaceCoefficients(bp);
        // draw surface
        drawSurface();
        // save the light correction surface
        if (generateLightingCorrectionSurfaceEyeCandy)
            saveFile("surface");
        // draw fixed image
        drawCorrectedImage(bp);
        // save the light corrected image
        saveFile("lightcorrected");
        return;
    }
    
    /**
     * Initializes the local order information for the least-squares fit.
     * Calculates the number of terms the n-th order polynomial
     * surface equation has.
     * 
     * <p>Pre: Assumes degree is a non-negative integer. If a negative order is passed, the order is calculated to be 0.
     * <br />Post: The order and orderTerms local members are initialized. The fit is ready to be calculated.  
     * @param degree The desired maximal order for the polynomial surface.
     * @return Returns the set order of polynomial surface.
     * @see #order 
     * @see #orderTerms 
     */
    public int setDegree (int degree) {
        // the maximum order in the polynomial surface
        if (degree >= 0) {
            order = degree; 
        } else {
            order = 0;
        }
        // the number of terms for a #order'd polynomial
        orderTerms = ((order+1)*(order+1)+order+1)/2; // "always" an even number (when it matters)
        return order;
    }

    /**
     * Based on the order, this method goes through the entire
     * process of calculating an A, b, and p matrix
     *  
     * <p>Pre: The setDegree() method was called and
     * <br />Post: The local member coefficients is initialized and prepared for use. The surface has been approximated. 
     * @param bp The imageProcessor for the original 8-bit grayscale image.
     * @see #setDegree
     */
    public void calculateSurfaceCoefficients(ImageProcessor bp) {
        int count=0, x=0, y=0, i=0, j=0, k=0, color=0, termCount=0;
        int xcount, ycount, multipleXs, multipleYs;
        int cutoff = 0;
        // if lightingCorrectionCutoff is negative, let it be based off the histogram of the image
        if (lightingCorrectionCutoff < 0)
            cutoff = getCutoffViaHistographThreshold(bp);
        else
            cutoff = lightingCorrectionCutoff;
        double [][] pixel;
        double [][] value;
        for (y=0; y<height; y++) {
            for (x=0; x<width; x++) {
                if (bp.getPixel(x,y) > cutoff) {
                    count ++;
                }
            }
        }
        if (displayDebugText)
            System.out.print("Found count to be:" + count + "\n");
        pixel = new double[count][orderTerms];
        value = new double[count][1];
        
        // Builds the A and b matrices composed of terms such as:
        // a,  b*x  ,  c*y  ,  d*x*x   ,  e*x*y  ,  f*y*y ,  .. etc
        count = 0;
        for (y=0; y<height; y++) {
            for (x=0; x<width; x++) {
                if (bp.getPixel(x,y) > cutoff) {
                    // generate an (order)th order surface polynomial across all permutations of x and y terms
                    termCount = 0;                  
                    for (i=0; i<=order; i++) {
                        for (j=0; j<=i; j++) {
                            pixel[count][termCount] = 1.0;
                            // a nice way to permute all the possibilities of x's and y's in each term
                            for (k=0;k<j; k++) {
                                pixel[count][termCount] *= (double)x;
                            }
                            for (k=j+1;k<=i; k++) {
                                pixel[count][termCount] *= (double)y;
                            }
                            termCount++;
                        }
                    }
                    // and the solution = the color
                    value[count][0] = (double)(bp.getPixel(x,y));
                    count ++;
                }
            }
        }
        if (displayDebugText)
            System.out.print("Made it out of the array loading faze.. creating matrices.\n");
        // perform least squares
        // load equation into matrix A
        // ax^2 + bxy + cy^2 + d + ex + fy= value
        Matrix A = new Matrix (pixel);
        if (displayDebugText)
            System.out.print("Created A matrix.\n");
        Matrix b = new Matrix (value);
        if (displayDebugText)
            System.out.print("Created b matrix.\n");
        Matrix p = getLeastSquaresProjection(A, b);      
        if (displayDebugText)
            System.out.print("Created p matrix.\n");
        coefficients = new double[orderTerms];
        for (i=0; i<orderTerms; i++) {
            coefficients[i] = p.get(i,0);
        }
    }
    
    /**
    * Performs a least-squares regression. This will return a matrix of the size (columns of A, 1).
    * Solves for the solution by the formula: p = (A^T * A)^-1 * A^T * b.
    * Where A , b is the function value at those datapoints,
    * and p is the desired projection vector that is the closest fit.
    *
    * <p>Pre: Assumes matrix A and B are properly formed and of appropriate dimension.
    * <br />Post: Returns the projection matrix (the solution) of the least-squares regression for matrices A and b. 
    * @param A The matrix of datapoints. Each row should be of the form x*x,x*y,y*y,1
    * @param b The matrix of function (color) values.
    * @return Returns a matrix (columns of A)x1 matrix that represent the constant values of the best fit projection.
    * @see #getCutoffViaHistographThreshold
    * @see #polynomialSurface
    */ 
    private Matrix getLeastSquaresProjection(Matrix A, Matrix b) {
        if (displayDebugText)
            System.out.print("line 1 - A rows="+A.getRowDimension()+"  columns="+A.getColumnDimension()+"\n");
        Matrix result = A.transpose();
        if (displayDebugText)
            System.out.print("line 2 - result rows="+result.getRowDimension()+"  columns="+result.getColumnDimension()+"\n");
        result = result.times(A);
        if (displayDebugText)
            System.out.print("line 3 - times A result rows="+result.getRowDimension()+"  columns="+result.getColumnDimension()+"\n");
        result = result.inverse();
        if (displayDebugText)
            System.out.print("line 4 - inverse result rows="+result.getRowDimension()+"  columns="+result.getColumnDimension()+"\n");
        result = result.times(A.transpose());
        if (displayDebugText) {
            System.out.print("line 5 - result rows="+result.getRowDimension()+"  columns="+result.getColumnDimension()+"\n");
            System.out.print("line 6 - b rows="+b.getRowDimension()+"  columns="+b.getColumnDimension()+"\n");
        }
        result = result.times(b);
        if (displayDebugText)
            System.out.print("line 7\n");
        return result;
    }

    /**
     * Analyzes the histograph for the most common colors.
     * I'm thinking I'll do an epsilon-neighborhood thing on the colors and see which is the largest.
     * Then I'll set that color (or that color+epsilon radius) to be the threshold cutoff.
     * 
     * Why do this? My motivation is I have all these images with black borders. Surface Plot 3D
     * gives a really nice quadratic surface except where it hits the bottom 0 value. This is to
     * stop those data points from damaging the quality of the least-squares fit later.
     * 
     * <p>Pre: The ImageProcessor is a grayscale image.
     * <br />Post: The method returns a computed threshold for pixels to be considered for lighting correction. 
     * @param bp The ImageProcessor for the image to allow for access to the histograph information.
     * @return Returns the value to ignore. Such that all colors of 0 to (that value) will be ignored for least-squares fitting.  
     */
    private int getCutoffViaHistographThreshold(ImageProcessor bp) {
        if (bp == null) {
            return 0;
        }
        ImageStatistics stats=imp.getStatistics();
        // assume the image is flat, select a small (dark) cutoff
        int cutoff = (int) (stats.mean - stats.stdDev);
        // if the image is from a reducing focal lens...
        if (stats.stdDev > stats.mean) {
            cutoff = (int)(stats.mean/4);
        }
        return cutoff;
    }
    
    /**
     * Returns the function value of the polynomial surface at the coordinate (x,y)
     * given the surface function f(x,y) = a + b*x + c*y + d*x^2 + ...
     * The polynomial surface is n-th order, where 'n' is the member variable 'order.'
     * The number of terms in the function is defined by orderTerms
     * 
     * <p>Pre: The given double[] eq array is not null.
     * <br />Post: The method returns the function value f(x,y) is returned as defined by the coefficients in eq.
     * @param x The x coordinate of the function.
     * @param y The y coordinate of the function.
     * @return Returns the result of the function f(x,y).
     * @see #getLeastSquaresProjection
     * @see #calculateSurfaceCoefficients
     */
    private int polynomialSurface(double x, double y) {
        int xcount, ycount, i, j, k, multipleXs, multipleYs;
        double total = 0, term;
        int termCount = 0;                  
        for (i=0; i<=order; i++) {
            for (j=0; j<=i; j++) {
                term = coefficients[termCount];
                // a nice way to permute all the possibilities of x's and y's in each term
                for (k=0;k<j; k++) {
                    term *= x;
                }
                for (k=j+1;k<=i; k++) {
                    term *= y;
                }
                total += term;
                termCount++;
            }
        }
        return (int)total;
    }
    
    /**
     * Draws the approximated surface onto a newly created image.
     * 
     * <p>Pre: The fit has already been performed. This requires the members order and orderTerms to be initialized. The passed double[] array must be the solution to the regression.
     * <br />Post: A new image is created and displayed.
     * @return Returns true on successful image creation and drawing. False otherwise.
     * @see #polynomialSurface
     * @see #calculateSurfaceCoefficients
     */
    public boolean drawSurface() {      
        if ((width < 1) || (height < 1) || (coefficients == null)) {
            return false;
        }

        String title = "Surface";
        ImagePlus surfaceIMP = NewImage.createByteImage (title, width, height, 1, NewImage.FILL_WHITE);
        ImageProcessor surfaceNP = surfaceIMP.getProcessor();
        int color, x, y;

        for (y=0; y<height; y++) {
            for (x=0; x<width; x++) {
                color = polynomialSurface((double)x,(double)y);
                if (color > 255) {
                    color = 255;
                } else if (color < 0) {
                    color = 0;
                }
                surfaceNP.putPixel(x,y,color);
            }
        }
        surfaceIMP.show();
        IJ.selectWindow(title);
        return true;
    }
    
    /**
     * Draws the light-corrected image onto a newly created image.
     * 
     * <p>Pre: The fit has already been performed. This requires the members order and orderTerms to be initialized. The passed double[] array must be the solution to the regression.
     * <br />Post: A new image is created and displayed.
     * @param bp The Image Processor for the original grayscale image.
     * @return Returns true on successful image creation and drawing. False otherwise.
     * @see #polynomialSurface
     * @see #calculateSurfaceCoefficients
     */
    public boolean drawCorrectedImage(ImageProcessor bp) {      
        if ((width < 1) || (height < 1) || (coefficients == null) || (bp == null)) {
            return false;
        }

        String title = "LightCorrected";
        ImagePlus imp = NewImage.createByteImage (title, width, height, 1, NewImage.FILL_WHITE);
        ImageProcessor np = imp.getProcessor();
        int color, x, y;
        
        for (y=0; y<height; y++) {
            for (x=0; x<width; x++) {
                color = polynomialSurface((double)x,(double)y);
                if (color > 255) {
                    color = 255;
                } else if (color < 0) {
                    color = 0;
                }
                color = bp.getPixel(x,y) - (color - 127); 
                if (color > 255) {
                    color = 255;
                } else if (color < 0) {
                    color = 0;
                }
                np.putPixel(x,y,color);
            }
        }
        imp.show();
        IJ.selectWindow(title);
        return true;
    }
}