/*

  javac -classpath .;graphics3d.jar *.java
  java -classpath .;graphics3d.jar Encoder filename

 */
import javax.swing.*;
import com.graphics3d.book.image.*;
import com.graphics3d.book.math.*;
import java.awt.*;

/**
   Takes the name of an image file as input and produces an Evolver
   program that produces a similar image.
   
   @author Morgan McGuire, <A
   HREF="mailto:morgan@cs.brown.edu">morgan@cs.brown.edu</A> */
public class Encoder {
    
    /**
       The window for display of the encoding process.
     */
    JFrame frame;

    /**
       Display of individual segment matches.
     */
    JPanel segmentDisplay;

    /**
       Display of individual segment matches.
     */
    JPanel segmentTargetDisplay;

    /**
       Display of the goal.
     */
    JPanel goalDisplay;

    JPanel infoDisplay;

    /**
       The input image
     */
    JImage input;

    /**
       Time spent evolving each segment.
     */
    int segmentTimeMinutes = 20;
    int finalTimeMinutes = 30;

    public static void main(String arg[]) {
        String filename = "", size = "";        

        if (arg.length == 2) {
            filename = arg[0];
            size = arg[1];
        } else {
            printHelp();
            System.exit(0);
        }
        
        filename = "../TestImages/test" + filename + "-" + size + ".jpg";

        Encoder encoder = new Encoder();
        encoder.run(filename);
    }

    private void buildGUI(String filename) {
        frame = new JFrame();

        frame.setTitle("Encoding " + filename);
        frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);

        JLabel l1 = new JLabel("Encoded");
        JLabel l2 = new JLabel("Original");
        l1.setPreferredSize(new Dimension(50, 20));
        l2.setPreferredSize(new Dimension(50, 20));
        
        segmentDisplay = new JPanel();
        segmentDisplay.setLayout(new BoxLayout(segmentDisplay, BoxLayout.X_AXIS));
        segmentDisplay.add(l1);
        segmentTargetDisplay = new JPanel();
        segmentTargetDisplay.setLayout(new BoxLayout(segmentTargetDisplay, BoxLayout.X_AXIS));
        segmentTargetDisplay.add(l2);
        goalDisplay = new JPanel();
        infoDisplay = new JPanel();
        
        JPanel all = new JPanel();
        frame.setContentPane(all);
        all.setLayout(new BoxLayout(all, BoxLayout.Y_AXIS));
        all.add(combine(goalDisplay, infoDisplay));
        all.add(segmentTargetDisplay);
        all.add(segmentDisplay);

        frame.pack();
        frame.show();
    }

    public void run(String filename) {

        // Initialize the Evolver system without the GUI
        System.out.println();
        GeneticParameters gp = new GeneticParameters("./");

        input = loadImage(filename);

        buildGUI(filename);

        System.out.println("Input: " + filename);

        Segmentation segmentation = new Segmentation(input, gp);
        Collage.setSegmentation(segmentation);

        // Show the goal image
        goalDisplay.add(caption(input, "Original"));
        //goalDisplay.add(JImageUtils.edgeImage(input));
        goalDisplay.add(caption(segmentation.getSegmentationImage(), "Segmentation"));
        /*
        goalDisplay.add(caption(JImageUtils.grayscale(JImageUtils.getFrequencies(input),
                                                      input.getPixelWidth(), input.getPixelHeight(), 3),
                                "Frequencies"));
        */

        {
            JTextArea t = new JTextArea(filename + "\n" + input.getPixelWidth() + 
                                        " x " + input.getPixelHeight() + "pixels");
            t.setBackground(Color.gray);
            goalDisplay.add(t);
        }
        frame.pack();
        

        ELExpression subPart[] = new ELExpression[segmentation.getNumSegments()];

        for (int s = 0; s < segmentation.getNumSegments(); s++) { 
            System.out.println();
            System.out.print("Matching segment #" + s + ".\nCreating new generation...");

            double minimumQuality = gp.minimumSegmentQuality;
            int timeout = segmentTimeMinutes * 60;
            MatchQualityFunction matcher = PerceptualMatchFunction.instance;
            boolean mask[] = segmentation.getMask(input.getPixelWidth(), input.getPixelHeight(), s);

            ELExpression start[] = expressionByFrequencyDecomposition(input, mask);

            start = null;
            subPart[s] = findSegmentMatch(mask, gp, input, matcher,
                                          minimumQuality, timeout, start, 0.50).expression;
        }

        Behavior result;

        // Create the collage
        {
            System.out.println("\n\nMatching whole image\n");
            ELExpression collage = new ELCall((ELOperator)gp.operatorsByName.get("Collage"), 
                                              subPart);
            double minimumQuality = 1.0;
            int timeout = finalTimeMinutes * 60;
            MatchQualityFunction matcher = PerceptualMatchFunction.instance;
        
            // Construct a mask for the whole image
            boolean mask[] = new boolean[input.getPixelWidth() * input.getPixelHeight()];
            for (int m = mask.length - 1; m >= 0; m--) {
                mask[m] = true;
            }

            // Now breed the collage
            result = findSegmentMatch(mask, gp, input, matcher,
                                      minimumQuality, timeout, new ELExpression[]{collage}, 0.90);
        }

        
        System.out.println("Done.");

        // TODO: save the picture, expression, and all data.
        
    }

    /**
       Finds an expression that the matcher returns greater than
       minimumQuality on against the indicated segment of the goal
       image.

       @param mask    Array indicating which areas to match
       @param timeout number of seconds to run before giving up
       @param start   Expressions to seed the generation with
       @param useStartProbability The probability that a given starting
                      slot will come from start and not a random behavior.
     */
    private Behavior findSegmentMatch(boolean mask[],
                                      GeneticParameters gp, JImage goalImage,
                                      MatchQualityFunction matcher,
                                      double minimumQuality, int timeout,
                                      ELExpression start[],
                                      double useStartProbability) {
        int width = goalImage.getPixelWidth(), height = goalImage.getPixelHeight();
        // Bugs elsewhere in the program tend to end up as zero size images.
        if ((width == 0) || (height == 0)) {
            throw new IllegalArgumentException("Width and height must be non-zero");
        }

        // Mask and display the goal segment for display purposes
        segmentTargetDisplay.add(JImageUtils.mask(goalImage, mask));
        frame.pack();
        
        long time0 = System.currentTimeMillis();
        double maxQuality = 0;

        int numBreeders = gp.numBreeders;
        int generationSize = gp.generationSize;
        
        // Make a new random generation (or all the same, in some cases)
        Behavior behavior[] = new Behavior[generationSize];
       
        for (int b = 0; b < generationSize; b++) {
            if ((start == null) || (Math.random() > useStartProbability)) {
                behavior[b] = new Behavior(gp, width, height);
            } else {
                // Choose a random one of our seed expressions
                int i = ((int)(Math.random() * start.length)) % start.length;
                behavior[b] = new Behavior(start[i], width, height, 0);
            }
        
            JImage output = behavior[b].toJImage();
            double quality = matcher.matchQuality(goalImage, output, mask);
            behavior[b].matchQuality = quality;
        }

        // Choose the best
        QuickSort.sort(behavior);
        Behavior parents[] = new Behavior[numBreeders];
        System.arraycopy(behavior, generationSize - numBreeders, parents, 0, numBreeders);
        behavior = parents;

        int gen = 0;
        infoDisplay.removeAll();
        infoDisplay.setLayout(new BoxLayout(infoDisplay, BoxLayout.Y_AXIS));
        while ((maxQuality < minimumQuality) && (System.currentTimeMillis() - time0 < timeout * 1000)) {
            maxQuality = 0;

            System.out.print("\nBreeding...");
            behavior = breed(goalImage, mask, parents, gp, generationSize, matcher);
            System.out.println("Done.\n\n");

            // Choose the best
            System.arraycopy(behavior, generationSize - numBreeders, parents, 0, numBreeders);            
            behavior = parents;

            gen++;

            // Evaluate the results and show off the best.
            for (int b = behavior.length - 1; b >= 0; b--) {
                JImage output = behavior[b].toJImage();
                double quality = matcher.matchQuality(goalImage, output, mask);

                if (quality > maxQuality) {
                    maxQuality = quality;
                }

                if (b == behavior.length - 1) {
                    // Show this result.
                    System.out.println(behavior[b]);
                    infoDisplay.removeAll();

                    // Show the color channel
                    infoDisplay.add(
                                    valign(
                                           halign(
                                                  caption(JImageUtils.mask(output, mask), "Comp. Color"),
                                                  caption(freqImage(output, mask), "Comp. Freq.")),
                                           halign(
                                                  caption(JImageUtils.mask(goalImage, mask), "Orig. Color"),
                                                  caption(freqImage(goalImage, mask), "Orig. Freq."))));
                    infoDisplay.add(new JTextArea("Quality = " + quality + "\nGeneration = " + gen));
                    frame.pack();
                }
            }
            
        }

        Behavior best = behavior[behavior.length - 1];
        segmentDisplay.add(JImageUtils.mask(best.toJImage(), mask));
        frame.pack();

        // Return the best
        return best;
    }

    /**
       Breeds the parents randomly, creating the desired number of children.
       The results are returned sorted from lowest to highest quality match.
     */
    static private Behavior[] breed(JImage goal, boolean mask[], 
                                    Behavior parent[], 
                                    GeneticParameters gp,
                                    int numChildrenDesired, 
                                    MatchQualityFunction matcher) {

        int width = goal.getPixelWidth(), height = goal.getPixelHeight();

        final int numRandom = numChildrenDesired / 15;
        int numParents = parent.length;
        Behavior child[] = new Behavior[numChildrenDesired + numParents + numRandom];

        int numChildren = 0;

        // Always breed against some random ones as well.
        {
            Behavior temp[] = new Behavior[(int)(parent.length * 1.10)];
            System.arraycopy(parent, 0, temp, 0, parent.length);
            for (int i = parent.length; i < temp.length; i++) {
                temp[i] = new Behavior(gp, width, height);
                temp[i].matchQuality = matcher.matchQuality(goal, temp[i].toJImage(), mask);
            }
            parent = temp;
        }

	while (numChildren < numChildrenDesired) {
	    System.out.print(".");
	    // Choose two parents (ok if they are the same)
	    int parentIndex1 = (int)Math.round((Math.random() * numParents));
	    int parentIndex2 = (int)Math.round((Math.random() * numParents));
            if (parentIndex1 == numParents) {
                parentIndex1 = numParents - 1;
            }
            if (parentIndex2 == numParents) {
                parentIndex2 = numParents - 1;
            }

	    Behavior parent1 = parent[parentIndex1];
	    Behavior parent2 = parent[parentIndex2];

	    ELExpression parentExpression1 = parent1.expression;
	    ELExpression parentExpression2 = parent2.expression;

	    // Breed
	    ELExpression childExpression = gp.mate(parentExpression1, parentExpression2);

	    // Evaluate the fitness of the child
	    Behavior childBehavior = 
		new Behavior(childExpression, width, height,
			     (int)Math.max(parent1.numGenerations, parent2.numGenerations));
      
	    boolean boring = false;

	    // Compare this child to other children so far.  It is boring if
	    // it is the same as a child we already have.
	    for(int c = 0; c < numChildren; c++) {
		if (childBehavior.sameExpression(child[c])) {
		    boring = true;
		    break;
		}
	    }
      
	    boring = boring || childBehavior.isBoring();

	    // Add this child if it is not boring.
	    if (! boring) {
                childBehavior.matchQuality = matcher.matchQuality(goal, childBehavior.toJImage(), mask);
		child[numChildren] = childBehavior;
		numChildren++;
	    }
	}

        // Add the parents back into the mix
        System.arraycopy(parent, 0, child, numChildren, numParents);

        // Add some random children
        numChildren += numParents;
        for (int i = 0; i < numRandom; i++) {
            child[numChildren] = new Behavior(gp, width, height);
            child[numChildren].matchQuality = matcher.matchQuality(goal, child[numChildren].toJImage(), mask);
            numChildren++;
        }

        QuickSort.sort(child);

        // Choose the best numChildrenDesired of these.
        Behavior result[] = new Behavior[numChildrenDesired];

        System.arraycopy(child, child.length - result.length, result, 0, numChildrenDesired);

        return result;
    }

    
    static void show(JComponent x) {
        show(x, "(Untitled)");
    }

    /**
       Opens a new JFrame containing x.  The frame will close the
       application when it is closed.  */
    static void show(JComponent x, String title) {
        JFrame frame = new JFrame();
        frame.setTitle(title);
        frame.setContentPane(x);
        frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        frame.pack();
        frame.show();
    }

    /**
      Creates a container for x and y.
     */
    static JPanel combine(JComponent x, JComponent y) {
        JPanel p = new JPanel();
        p.add(x);
        p.add(y);
        return p;
    }

    static JPanel combine(JComponent x, JComponent y, JComponent z) {
        JPanel p = new JPanel();
        p.add(x);
        p.add(y);
        p.add(z);
        return p;
    }

    /**
       Prints help information to the command line.
     */
    static void printHelp() {
        // TODO
    }

    /**
       Attempts to load an image, exiting the program if anything
       goes wrong.
     */
    static JImage loadImage(String filename) {
        JImage image = null;

        try {
            image = new JImage(filename);
        } catch (java.io.IOException e) {
            System.err.println("'" + filename + "' not found.");
            System.exit(1);
        } catch (com.graphics3d.book.image.ImageLoadingException e) {
            System.err.println("'" + filename + "' corrupted.");
            System.exit(1);
        }

        return image;
    }


    static private JComponent halign(JComponent a, JComponent b) {
        JPanel panel = new JPanel();
        panel.setLayout(new BoxLayout(panel, BoxLayout.X_AXIS));
        panel.add(a);
        panel.add(b);
        return panel;
    }

    static private JComponent valign(JComponent a, JComponent b) {
        JPanel panel = new JPanel();
        panel.setLayout(new BoxLayout(panel, BoxLayout.Y_AXIS));
        panel.add(a);
        panel.add(b);
        return panel;
    }

    static private JComponent caption(JComponent a, String caption) {
        return valign(a, new JLabel(caption));
    }

    static private JImage freqImage(JImage a, boolean mask[]) {
        return JImageUtils.mask(JImageUtils.grayscale(JImageUtils.getFrequencies(a),
                                                      a.getPixelWidth(),
                                                      a.getPixelHeight(), 3), mask);
    }


    static private ELVector findAverageColor(JImage a, boolean mask[]) {
        float avg[] = JImageUtils.averageColor(a, mask);
        return new ELVector(avg[0], avg[1], avg[2]);
    }

    /**
       Computes the average color of a within the masked region and
       returns it in outColor.  The result JImage is a new JImage with
       outColor subtracted from all elements.
    */
    static private JImage subtractAverage(JImage a, boolean mask[], float outColor[]) {
        if (outColor.length != 3) {
            throw new IllegalArgumentException("outColor must have length 3");
        }

        double red[] = a.getChannel(JImage.RED);
        double grn[] = a.getChannel(JImage.GREEN);
        double blu[] = a.getChannel(JImage.BLUE);

        // Compute the average color
        outColor[0] = (float)JImageUtils.avg(red, mask);
        outColor[1] = (float)JImageUtils.avg(grn, mask);
        outColor[2] = (float)JImageUtils.avg(blu, mask);

        // Subtract
        JImage out = new JImage(a.getPixelWidth(), a.getPixelHeight());
        for (int i = red.length - 1; i >= 0; i--) {
            red[i] -= outColor[0];
            grn[i] -= outColor[1];
            blu[i] -= outColor[2];
        }

        out.setChannel(JImage.RED,   red);
        out.setChannel(JImage.GREEN, grn);
        out.setChannel(JImage.BLUE,  blu);

        return out;
    }

    /**
       Directly computes a reasonable expression
     */
    private ELExpression[] expressionByFrequencyDecomposition(JImage input, boolean mask[]) {
        ELExpression freq0 = findAverageColor(input, mask);
        ELExpression freq1 = new ELCall(new Noise(), new ELExpression[]{});
        ELExpression c1_0 = new ELCall(new Add(), new ELExpression[]{freq0, freq1});


        return new ELExpression[]{freq0, freq1, c1_0};
    }
}

