
import com.graphics3d.book.image.*;

/**
   Routines for working with JImages.
 */
abstract public class JImageUtils {

    /**
       Produces a new image that is black everywhere ma is false and
       elsewhere draws its pixels from image.
    */
    static public JImage mask(JImage image, boolean ma[]) {
        JImage out = new JImage(image.getPixelWidth(), image.getPixelHeight());

        int pixel[] = out.getPixels();
        int in[] = image.getPixels();
        
        if (pixel.length != ma.length) {
            throw new IllegalArgumentException("Wrong dimensions");
        }
        
        for (int p = pixel.length - 1; p >= 0; p--) {
            pixel[p] = ma[p] ? in[p] : 0x00000000;
        }

        return out;
    }


    /**
       Returns a measure of how well the pixel intensity values match
       between the two images.  RMS pixel error is used and scaled to
       0..1.  Intensity is measured as (R + G + B) / 3.
       
       <PRE>
       1.0 = perfect match
       0.0 = completely off (one is white, the other is black)
       </PRE> */
    static public double intensityMatch(JImage a, JImage b, boolean mask[]) {
        checkDimensions(a, b);
        
        return 1.0 -  RMS(div(sub(add(add(a.getChannel(JImage.RED), a.getChannel(JImage.BLUE)),
                                      a.getChannel(JImage.GREEN)), 
                                  add(add(b.getChannel(JImage.RED), b.getChannel(JImage.BLUE)),
                                      b.getChannel(JImage.GREEN))), 3),
                          mask);
    }

    /**
       Returns a measure of how well the edges match between the two
       images.  Average edge match error is used and scaled to 0..1.
       
       <PRE> 1.0 = perfect match
       0.0 = completely off
       </PRE> */
    static public double edgeMatch(JImage a, JImage b) {
        checkDimensions(a, b);

        final boolean edgeA[] = getEdges(a);
        final boolean edgeB[] = getEdges(b);
        return edgeMatch(edgeA, edgeB);
    }

    static public double edgeMatch(boolean edgeA[], boolean edgeB[]) {
        double sum = 0;
        for (int i = edgeA.length - 1; i >= 0; i--) {
            sum += (edgeA[i] == edgeB[i]) ? 1 : 0;
        }
        
        return sum / edgeA.length;
    }

    /**
       Both arguments must be absolute gradients.
     */
    static public double frequencyMatch(double a[], double b[], boolean mask[]) {
        return RMS(sub(a, b), mask);
    }

    /**
       Produces an image of the same dimensions as a that is black
       where there are edges in a and white elsewhere.
     */
    static public JImage edgeImage(JImage a) {
        final boolean edgeA[] = getEdges(a);
        final int black = JImage.createPixel(0,0,0);
        final int white = JImage.createPixel(1,1,1);
        
        final JImage result = new JImage(a.getPixelWidth(), a.getPixelHeight());
        
        final int pixels[] = result.getPixels();

        for (int i = pixels.length - 1; i >= 0; i--) {
            pixels[i] = edgeA[i] ? black : white;
        }
        
        result.flush();

        return result;
    }

    /**
       Approximates the highest frequency at each pixel.
     */
    static public double[] getFrequencies(JImage image) {
        final int width = image.getPixelWidth();
        final int height = image.getPixelHeight();
        final int length = width * height;
        final float grad[] = perceptualGradient(image);
        final double out[] = new double[length];

        final double weight[] = {0.00, 0.01, 0.02, 0.01, 0.00,
                                 0.01, 0.17, 0.25, 0.17, 0.01,
                                 0.02, 0.25, 1.00, 0.25, 0.02,
                                 0.01, 0.17, 0.25, 0.17, 0.01,
                                 0.00, 0.02, 0.02, 0.01, 0.00};
        
        // Normalize weights
        double sum = 0;
        for (int i = 0; i < 25; i++) {
            sum += weight[i];
        }
        for (int i = 0; i < 25; i++) {
            weight[i] /= sum;
        }
              

        // Blur
        for (int y = 2; y < height - 2; y++) {
            for (int x = 2; x < width - 2; x++) {
                int j = x + y * width;
                for (int dx = -2; dx <= +2; dx++) {
                    for (int dy = -2; dy <= +2; dy++) {
                        int i = (x + dx) + (y + dy) * width;
                        out[j] += grad[i] * weight[(dx + 2) + (dy + 2) * 5];
                    }
                }

                // log to get to some reasonably linear scaling
                out[j] = Math.log(out[j] * 99 + 1) / 10;
            }
        }

        return out;
    }

    /**
       Computes the average color over the masked region and
       returns it as an array of three floats.
     */
    static public float[] averageColor(JImage a, boolean mask[]) {
        double red[] = a.getChannel(JImage.RED);
        double grn[] = a.getChannel(JImage.GREEN);
        double blu[] = a.getChannel(JImage.BLUE);
        
        // Compute the average color
        return new float[]{(float)JImageUtils.avg(red, mask),
                           (float)JImageUtils.avg(grn, mask),
                           (float)JImageUtils.avg(blu, mask)};
    }

    static public float[] perceptualGradient(JImage image) {
        final int width = image.getPixelWidth();
        final int height = image.getPixelHeight();
        final int length = width * height;
        final float out[] = new float[length];

        // perceptual weighting of colors
        final double colorWeight[] = new double[4];
        colorWeight[JImage.ALPHA] = 0;
        colorWeight[JImage.RED]   = 0.299;
        colorWeight[JImage.GREEN] = 0.587;
        colorWeight[JImage.BLUE]  = 0.114;

        // RGB separation
        for (int channel = JImage.RED; channel < JImage.BLUE; channel++) {
            // Compute the magnitude of the gradient for this channel
            double value[] = image.getChannel(channel);
            double weight = colorWeight[channel];
            for (int y = 1; y < height; y++) {
                for (int x = 1; x < width; x++) {
                    int i0 = x + y * width;
                    int ix = (x - 1) + y * width;
                    int iy = x + (y - 1) * width;
                    double dx = value[ix] - value[i0];
                    double dy = value[iy] - value[i0];
                    
                    out[i0] += (float)((dx * dx + dy * dy) * (weight * weight));
                }
            }
        }
        
        // Normalize out
        for (int i = 0; i < length; i++) {
            out[i] = (float)Math.sqrt(out[i]);
        }

        return out;
    }

    static public JImage grayscale(double value[], int width, int height, double scale) {
        JImage out = new JImage(width, height);
        
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                double v = value[x + y * width] * scale;
                out.set(x, y, JImage.createPixel(v,v,v));
            }
        }

        return out;

    }

    /**
       Finds edges.  An edge is a location where any channel (r,g,b) changes by
       more than 10% between two horizontal or vertical pixels.  After finding,
       edges are smeared by 2 pixels in each direction to thicken them.
     */
    static public boolean[] getEdges(JImage image) {
        final int length = image.getPixelWidth() * image.getPixelHeight();
        final boolean result[] = new boolean[length];

        final double cutoff = 0.10;

        final int width = image.getPixelWidth();
        final int height = image.getPixelHeight();

        for (int c = JImage.RED; c <= JImage.BLUE; c++) {
            double pixel[] = image.getChannel(c);

            // find horizontal edges
            for (int y = 0; y < height; y++) {
                for (int x = 1; x < width; x++) {
                    int i = x + y * width;
                    result[i] = result[i] || (Math.abs(pixel[i] - pixel[i - 1]) > cutoff);
                }
            }

            // find horizontal edges
            for (int x = 0; x < width; x++) {
                for (int y = 1; y < height; y++) {
                    int i = x + y * width;
                    result[i] = result[i] || (Math.abs(pixel[i] - pixel[i - width]) > cutoff);
                }
            }
        }

        boolean smeared[] = new boolean[length];
        final int s = 2;
        for (int y = s; y < height - s; y++) {
            for (int x = s; x < width - s; x++) {
                int is = x + y * width;
                for (int dx = -s; dx <= s; dx++) {
                    for (int dy = -s; dy <= s; dy++) {
                        int ir = (x + dx) + (y + dy) * width;
                        smeared[is] = smeared[is] || result[ir];
                    }
                }
            }
        }
        
        return smeared;
    }
    
    /**
       Returns a measure of how well the pixel values match between
       the two images.  RMS pixel error is used and scaled to 0..1.
       
       Matching green is considered twice as important.

       <PRE>
       1.0 = perfect match
       0.0 = completely off (one is white, the other is black)
       </PRE>
    */
    static public double colorMatch(JImage a, JImage b, boolean mask[]) {
        // Compute the sum squares for each channel
        checkDimensions(a, b);
        
        double sum = 0;
        for (int c = JImage.RED; c <= JImage.BLUE; c++) {
            double x = RMS(sub(a.getChannel(c), b.getChannel(c)), mask);
            sum += x;
        }

        // Compute the mean.  Note that the mean error will be between 0 and 1, since
        // all of the values are in the range 0...1
        return 1.0 - (sum / 3.0);
    }

    static public JImage blur(JImage a) {
        int width = a.getPixelWidth();
        int height = a.getPixelHeight();
        int in[] = a.getPixels();
        JImage outImage = new JImage(width, height);
        int out[] = outImage.getPixels();

        final int blurWidth = (int)(Math.sqrt(in.length) / 15);
        final int blurHeight = blurWidth / 2;

        // Blur to eliminate high frequencies and help
        // our run length encoder out.

        // Horizontal
        for (int i = in.length - 1 ; i >= blurWidth; i--) {
            double sr = 0, sg = 0, sb = 0;

            for (int j = 0; j < blurWidth; j++) {
                double r = JImage.pixelRed(in[i - j]);
                double g = JImage.pixelGreen(in[i - j]);
                double b = JImage.pixelBlue(in[i -j ]);
                
                sr += r;
                sg += g;
                sb += b;
            }

            out[i] = JImage.createPixel(sr / blurWidth, sg / blurWidth, sb / blurWidth);
        }

        // Elements we missed
        for (int i = 0; i < blurWidth; i++) {
            out[i] = in[i];
        }

        // Vertical (doesn't wrap around like the horizontal one)
        for (int x = 0; x < width; x++) {
            for (int y = blurHeight; y < height; y++) {
                double sr = 0, sg = 0, sb = 0;
                int i = x + y * width;
                for (int j = 0; j < blurHeight; j++) {
                    int o = i - j * width;
                            
                    double r = JImage.pixelRed(in[o]);
                    double g = JImage.pixelGreen(in[o]);
                    double b = JImage.pixelBlue(in[o]);
                    
                    sr += r;
                    sg += g;
                    sb += b;
                }
                out[i] = JImage.createPixel(sr / blurHeight, sg / blurHeight, sb / blurHeight);
            }
        }

        // Stretch out the top rows to cover the fact that we didn't blur them
        for (int y = 0; y < blurHeight; y++) {
            System.arraycopy(out, width * blurHeight, out, y * width, width);
        }
        

        return outImage;
    }
    
    /**
       colorMatch on blurred versions of the images.
     */
    static public double localAverageColorMatch(JImage a, JImage b, boolean mask[]) {
        return colorMatch(blur(a), blur(b), mask);
    }

    /**
       Returns the RMS error between the average colors of a and b.
     */
    static public double averageColorMatch(JImage a, JImage b, boolean mask[]) {
        // Compute the sum squares for each channel
        checkDimensions(a, b);
        
        float ac[] = averageColor(a, mask);
        float ab[] = averageColor(b, mask);

        double sum = Math.sqrt(Math.pow(ac[0] - ab[0], 2) + 
                               Math.pow(ac[1] - ab[1], 2) + 
                               Math.pow(ac[2] - ab[2], 2));
        return 1.0 - (sum / 3.0);
    }


    /**
       Throws an exception if a and b do not have the same dimensions.
     */
    static private void checkDimensions(JImage a, JImage b) {
        if ((a.getPixelWidth() != b.getPixelWidth()) ||
            (a.getPixelHeight() != b.getPixelHeight())) {
            throw new IllegalArgumentException("Mismatched sizes.");
        }
    }

    /**
       Returns the difference (subtraction) x - y.
     */
    static private double[] sub(double x[], double y[]) {
        int i;
        double result[] = new double[x.length];

        if (x.length != y.length) {
            throw new IllegalArgumentException("Size mismatch");
        }

        for (i = x.length - 1; i >= 0; i--) {
            result[i] = x[i] - y[i];
        }

        return result;
    }

    /**
       Returns the quotient x / y.
     */
    static private double[] div(double x[], double y) {
        int i;
        double result[] = new double[x.length];

        for (i = x.length - 1; i >= 0; i--) {
            result[i] = x[i] / y;
        }

        return result;
    }


    /**
       Returns the sum x + y.
     */
    static private double[] add(double x[], double y[]) {
        int i;
        double result[] = new double[x.length];

        if (x.length != y.length) {
            throw new IllegalArgumentException("Size mismatch");
        }

        for (i = x.length - 1; i >= 0; i--) {
            result[i] = x[i] + y[i];
        }

        return result;
    }

    /**
       Computes the sum of the squares of x.
     */
    static private double sumSquares(double x[]) {
        double sum = 0;
        for (int i = x.length - 1; i >= 0; i--) {
            sum += x[i] * x[i];
        }
        return sum;
    }

    /**
       Computes the RMS x, taking the mask into account.
     */
    static private double RMS(double x[], boolean mask[]) {
        int count = 0;
        double sum = 0;
        for (int i = x.length - 1; i >= 0; i--) {
            if (mask[i]) {
                sum += x[i] * x[i];
                count++;
            }
        }
        return Math.sqrt(sum / count);
    }

    /**
       Computes the avg x, taking the mask into account.
     */
    static public double avg(double x[], boolean mask[]) {
        int count = 0;
        double sum = 0;
        for (int i = x.length - 1; i >= 0; i--) {
            if (mask[i]) {
                sum += x[i];
                count++;
            }
        }
        return sum / count;
    }

}
