import java.awt.Color;
import com.graphics3d.book.image.JImage;

/**
   Maps pixel positions to expressions used to generate those pixels.


   Segmentation uses A.R. Smith's RGB to HSV algorithm, then quantizes
   the color space into light, dark, and numBins - 2 color bins.

   @author Morgan McGuire, morgan@cs.brown.edu 
*/
public class Segmentation {

    /**
       Number of tones the segmenter should attempt to break the image
       into.  This should be a power of two.  */
    private int numBins;

    private int numHueBins;

    /**
       Number of bins actually in use.
     */
    private int numBinsUsed;
    public int getNumSegments() {
        return numBinsUsed;
    }

    /**
       binIndex[i] is the renumbered index of the ith bin.
       If it is -1, then that bin is not in use.  This is used
       to present an interface to the rest of the world of consecutively 
       numbered bins.
     */
    private int binIndex[];

    /**
       segmentIndex[i] is what internal bin i is called internally.
       This is the inverse of binIndex.
     */
    private int segmentIndex[];

    private ELExpression expression[];

    /**
       Width of the segmentation image.
     */
    private int width;

    /**
       Height of the segmentation image.
     */
    private int height;

    /**
       expression[segmentation[i]] colors pixel i Bin 0 is dark, bin 1
       is light, and bins 2..numBins-1 are colors.
    */
    private int segmentation[];

    /**
       If binUsed[i] is false, then expression[i] is null.
     */
    private boolean binUsed[];

    /**
       Expand this segmentation by a factor of scale in each dimension.
     */
    public void scaleBy(int scale) {
        int out[] = new int[width * height * scale * scale];

        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                int b = segmentation[x + y * width];
                for (int dx = 0; dx < scale; dx++) {
                    for (int dy = 0; dy < scale; dy++) {
                        out[(x * scale + dx)     + (y * scale + dy)     * (width * scale)] = b;
                    }
                }
            }
        }
        
        segmentation = out;
        width *= scale;
        height *= scale;
    }

    /**
       Creates a segmentation from the given image.
       The image is not modified.
     */
    public Segmentation(JImage image, GeneticParameters gp) {
        numBins = gp.numSpatialSegmentationBins;
        numHueBins = numBins - 2;

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

        segmentation = new int[width * height];
        int pixel[] = preProcess(image.getPixels(), width, height);
        
        // Histogram for bins
        int count[] = new int[numBins];

        for (int i = pixel.length - 1; i >= 0; i--) {
            int p = pixel[i];
            int bin = pixelToBin(p);
            
            count[bin]++;
            segmentation[i] = bin;
        }

        // Find bins with few elements and collapse them into other bins.

        // Minimum number of elements that must be in a bin for it to
        // be preserved.
        int cutoff = pixel.length / (numBins * 4);
        binUsed = new boolean[numBins];
        binIndex = new int[numBins];
        segmentIndex = new int[numBins];
        for (int bin = 0; bin < numBins; bin++) {
            binUsed[bin] = (count[bin] > cutoff);
            if (binUsed[bin]) {
                binIndex[bin] = numBinsUsed;
                segmentIndex[numBinsUsed] = bin;
                numBinsUsed++;
            } else {
                binIndex[bin] = -1;
            }
            System.out.println("" + count[bin] + " - " + binUsed[bin]);
        }

        // remap[bin] is the bin where we put pixels that used to
        // contain bin.
        int remap[] = new int[numBins];
        for (int bin = 0; bin < numBins; bin++) {

            // Start with identity
            remap[bin] = bin;

            // Find another bin to put these pixels in
            while (!binUsed[remap[bin]]) {
                remap[bin] = (remap[bin] + 1) % numBins;
            }
        }

        // Now collapse little used bins
        for (int i = segmentation.length - 1; i >= 0; i--) {
            segmentation[i] = remap[segmentation[i]];
        }

        expression = new ELExpression[numBins];
    }

    /**
       Normalize the intensities of an image before segmentation.
       Returns a new array of pixels.
     */
    private int[] preProcess(int in[], int width, int height) {
        
        int out[] = new int[in.length];

        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);
        }

        /*
        // Find the average intensity of the image along each color
        // channel.
        double ar = 0;
        double ag = 0;
        double ab = 0;
        for (int i = in.length - 1; i >= 0; i--) {
            ar += JImage.pixelRed(out[i]);
            ag += JImage.pixelGreen(out[i]);
            ab += JImage.pixelBlue(out[i]);
        }

        ar /= in.length;
        ag /= in.length;
        ab /= in.length;
        
        // Rescale to move the average to 0.5 along each channel.
        for (int i = in.length - 1; i >= 0; i--) {
            double r = JImage.pixelRed(out[i]);
            double g = JImage.pixelGreen(out[i]);
            double b = JImage.pixelBlue(out[i]);

            if (r < ar) {
                r = 0.5 * r / ar;
            } else {
                r = (0.5 * (r - ar) / (1 - ar)) + 0.5;
            }

            if (g < ab) {
                g = 0.5 * g / ag;
            } else {
                g = (0.5 * (g - ag) / (1 - ag)) + 0.5;
            }

            if (b < ab) {
                b = 0.5 * b / ab;
            } else {
                b = (0.5 * (b - ab) / (1 - ab)) + 0.5;
            }

            out[i] = JImage.createPixel(r, g, b);
        }

        */        
        return out;

    }

    private static final float hsv[] = new float[3];
    /**
       Computes the number of the bin that this pixel should be classified in.
     */
    private int pixelToBin(int p) {
        float r = (float)JImage.pixelRed(p);
        float g = (float)JImage.pixelGreen(p);
        float b = (float)JImage.pixelBlue(p);

        // Move to HSB space
        Color.RGBtoHSB((int)(256 * r), (int)(256 * g), (int)(256 * b), hsv);

        float h = hsv[0], s = hsv[1], v = hsv[2];

        int bin = ((int)Math.round((hsv[0] + (1.0 / numHueBins)) * numHueBins)) % numHueBins;

        // light (v > .8) goes into bin 1 regardless of hue.
        if ((v > 0.8) && (s < 0.3)) {
            return 1;
        } else if (v < 0.3) {
            return 0;
        }
            
        return bin + 2;
    }


    /**
       Returns true if this segmentation is sampled at exactly these dimensions.
     */
    public boolean isSameSize(int width, int height) {
        return (this.width == width) && (this.height == height);
    }

    /**
       Gets the weight index that will have value 1.0 where i = x + y * width.
       Only call if isSameSize() returned true.
     */
    final public int getBin(int i) {
        return binIndex[segmentation[i]];
    }

    /**
       Returns an array of weights indicating how much each segment
       covers the given sample pixel, which is from a w x h grid.

       The indices on the weights are between 0 and numSegments - 1.
    */
    public float[] getWeights(int x1, int y1, int x2, int y2, int w, int h, float[] out) {
        if (out.length < numBinsUsed) {
            out = new float[numBinsUsed];
        }

        int i = (int)Math.round(x1 * (width - 1) / (double)(w - 1));
        int j = (int)Math.round(y1 * (height - 1) / (double)(h - 1));

        // Look up which segment colors this pixel, convert to the
        // external index, and set the weight corresponding to that
        // segment to 1.0.
        out[binIndex[segmentation[i + j * width]]] = 1.0f;

        // TODO: could do fractional weighting

        return out;
    }
    

    /**
       Produces a mask that is true for every pixel that belongs in
       segment and false elsewhere.  The mask is row major order and
       has dimensions width x height, which must match the original
       segmentation dimensions.
    */
    public boolean[] getMask(int width, int height, int segment) {
        if (! (isSameSize(width, height))) {
            throw new IllegalArgumentException("Wrong dimensions");
        }

        int bin = segmentIndex[segment];
        boolean out[] = new boolean[width * height];

        for (int i = out.length - 1; i >= 0; i--) {
            out[i] = (segmentation[i] == bin);
        }

        return out;
    }

    /**
       Gets the mask at scale*resolution
     */
    public boolean[] getMaskN(int width, int height, int scale, int segment) {
        int bin = segmentIndex[segment];
        boolean out[] = new boolean[width * height * scale * scale];

        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                boolean b = (segmentation[x + y * width] == bin);
                out[x * scale + (y * scale) * (width * scale)] = b;
                out[(x * scale + 1) + (y * scale) * (width * scale)] = b;
                out[(x * scale) + (y * scale + 1) * (width * scale)] = b;
                out[(x * scale + 1) + (y * scale + 1) * (width * scale)] = b;
            }
        }

        return out;
    }

    /**
       Produces a new JImage that is black everywhere that doesn't
       correspond to the indicated segment.

       image must have the same dimensions as the segmentation.
     */
    public JImage maskSegment(JImage image, int segment) {
        if (segment > getNumSegments()) {
            throw new IllegalArgumentException("Bad segment");
        }

        JImage out = new JImage(image.getPixelWidth(), image.getPixelHeight());
        int bin = segmentIndex[segment];

        int pixel[] = out.getPixels();
        int in[] = image.getPixels();
        
        if (! (isSameSize(image.getPixelWidth(), image.getPixelHeight()))) {
            throw new IllegalArgumentException("Wrong dimensions");
        }
        
        for (int p = pixel.length - 1; p >= 0; p--) {
            if (segmentation[p] != bin) {
                // Knock out the pixel
                pixel[p] = 0x0000000;
            } else {
                pixel[p] = in[p];
            }
        }

        return out;
    }

    private static final double rgb[] = new double[3];
    /**
       Computes a pixel value from bin index.
     */
    private int binToPixel(int bin) {

        float h,s,v;

        if (bin == 0) {
            // Dark bin
            h = 0.0f; 
            s = 1.0f;
            v = 0.0f;
        } else if (bin == 1) {
            // Light bin
            h = 0.0f;
            s = 0.0f;
            v = 1.0f;
        } else {
            // Color bin
            h = (bin - 2) / (float)numHueBins;
            s = 1.0f;
            v = 1.0f;
        }
        
        Color c = new Color(Color.HSBtoRGB(h, s, v));

        return JImage.createPixel(c.getRed() / 255.0, 
                                  c.getGreen() / 255.0, 
                                  c.getBlue() / 255.0);
    }

    /**
       Creates a picture of the segmentation.
     */
    public JImage getSegmentationImage() {
        // Compute a table mapping hues to pixel values so we 
        // don't recompute every pixel.
        final int binToPixel[] = new int[numBins];
        for (int i = 0; i < numBins; i++) {
            binToPixel[i] = binToPixel(i);
        }

        final JImage image = new JImage(width, height);
        final int pixel[] = image.getPixels();
        for (int i = pixel.length - 1; i >= 0; i--) {
            pixel[i] = binToPixel[segmentation[i]];
        }

        return image;
    }
}





