package ac.essex.gp.problems;

import ac.essex.gp.training.segmentation.TrainingImage;
import ac.essex.gp.training.segmentation.JasminePixelSelection;
import ac.essex.gp.params.GPParams;
import ac.essex.gp.params.NodeParams;
import ac.essex.gp.individuals.Individual;
import ac.essex.gp.util.DataStack;
import ac.essex.gp.nodes.*;
import ac.essex.gp.nodes.ercs.*;
import ac.essex.gp.nodes.imaging.*;
import ac.essex.gp.nodes.imaging.texture.Mean;
import ac.essex.gp.nodes.imaging.texture.Variance;
import ac.essex.gp.Evolve;
import ac.essex.ooechs.imaging.commons.apps.jasmine.JasmineProject;
import ac.essex.ooechs.imaging.commons.apps.jasmine.JasmineImage;
import ac.essex.ooechs.imaging.commons.apps.jasmine.JasmineClass;
import ac.essex.ooechs.imaging.commons.Pixel;

import java.util.Vector;
import java.io.File;
import java.awt.image.BufferedImage;
import java.awt.*;

/**
 * <p/>
 * 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,
 * provided that any use properly credits the author.
 * 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 at http://www.gnu.org
 * </p>
 *
 * @author Olly Oechsle, University of Essex, Date: 18-Jan-2007
 * @version 1.0
 */

public class JasmineSegmentationProblem extends Problem {

    public String getName() {
        return "Segmentation Problem";
    }

    // the training data
    public Vector<ac.essex.gp.training.segmentation.TrainingImage> trainingData;

    // the project that we're using
    JasmineProject project;


    public JasmineSegmentationProblem(File jasmineFile) throws Exception {
        // load the jasmine project
        project = JasmineProject.load(jasmineFile);
    }

    public JasmineSegmentationProblem(JasmineProject p) {
        project = p;
    }

    public void initialise(Evolve e, GPParams params) {
        // produce the training data
        trainingData = new Vector<TrainingImage>(10);

        try {

            // now get its images
            for (int i = 0; i < project.getImages().size(); i++) {

                JasmineImage image = project.getImages().elementAt(i);

                BufferedImage img = image.getBufferedImage();

                try {

                    ac.essex.gp.training.segmentation.PixelSelection selection = new JasminePixelSelection(image);

                    if (selection.getPixels().size() > 0) {

                        trainingData.add(new TrainingImage(img, selection));

                    }

                } catch (Exception err) {                    
                    //System.err.println("Cant get pixel selection for image: " + i);
                }

            }

            if (trainingData.size() == 0) {
                e.fatal("No overlays found - GP cannot proceed without training data.");
            }

            // set up the classes that return is allowed to ... return
            Return.classes = new int[project.getPixelClasses().size()];
            for (int i = 0; i < project.getPixelClasses().size(); i++) {
                Return.classes[i] = project.getPixelClasses().elementAt(i).classID;
            }

        } catch (Exception err) {
            e.fatal("GP system cannot load Jasmine project:\n" + err.toString());
            err.printStackTrace();
        }

        // initialise the parameters
        // register the nodes we want to use
        params.registerNode(new Add());
        params.registerNode(new Mul());
        params.registerNode(new Sub());

        params.registerNode(new If());
        params.registerNode(new ac.essex.gp.nodes.logic.More());
        params.registerNode(new ac.essex.gp.nodes.logic.Less());

        params.registerNode(new Return());

        params.registerNode(new Red());
        params.registerNode(new Green());
        params.registerNode(new Blue());
        params.registerNode(new Hue());
        params.registerNode(new Saturation());
        params.registerNode(new Lightness());

        params.registerNode(new Variance());
        params.registerNode(new Mean());

        //params.registerNode(new Range());

        //params.registerNode(new HorizontalSobel());
        //params.registerNode(new VerticalSobel());

        params.registerNode(new LargeIntERC());
        params.registerNode(new TinyDoubleERC());
        params.registerNode(new SmallIntERC());
        params.registerNode(new BoolERC());

        // set up additional parameters
        params.setReturnType(NodeParams.SUBSTATEMENT);
        //params.setGenerations(500);

    }

    public void customiseParameters(GPParams params) {}

    public int getClassCount() {
        return project.getPixelClasses().size();
    }

    public void evaluate(Individual ind, DataStack data) {

        int fitness = 0;
        int hits = 0;

        boolean[] classHits = new boolean[project.getPixelClasses().size() + 1];

        for (int i = 0; i < trainingData.size(); i++) {
            TrainingImage image = trainingData.elementAt(i);

            // now go through the data on the image (but only look at the selected pixels)
            ac.essex.gp.training.segmentation.PixelSelection selection = image.getSelection();

            // set up the image on the stack
            data.setImage(image);

            int x = 0;

            Vector<Pixel> pixels = selection.getPixels();

            for (int j = 0; j < pixels.size(); j++) {

                Pixel point = pixels.elementAt(j);

                x++;

                // speeds up execution and tends to get smoother results
                if (x % 2 == 0) continue;

                // save this to the data stack. The nodes are able to access this stack.
                data.setX(point.x);
                data.setY(point.y);

                // runt the individual
                int result = (int) ind.execute(data);

                if (!data.usesImaging) {
                    // don't execute the individual many times if it doesn't even use imaging
                    break;
                }

                // record that this class was hit (even if it was a FP)
                if (result > 0 && result < classHits.length) classHits[result] = true;

                // wrong answers are penalised
                if (result != point.value) {
                    fitness++;
                } else {
                    hits++;
                }

            }

        }

        // award bonus of 15% for each class correctly hit
        for (boolean classHit : classHits) {
            if (classHit) fitness *= 0.85;
        }

        if (!data.usesImaging) {
            fitness = Integer.MAX_VALUE;
        }

        ind.setKozaFitness(fitness);
        ind.setHits(hits);

    }

    public BufferedImage describe(Individual ind, DataStack data, int index) {

        // make a cache of the colours so we don't have to requery the vector all the time
        Color[] colors = new Color[project.getShapeClasses().size() + project.getPixelClasses().size() + 1];
        for (int i = 0; i < project.getPixelClasses().size(); i++) {
            JasmineClass jasmineClass = project.getPixelClasses().elementAt(i);
            colors[i + 1] = jasmineClass.color;
        }

        // get the first training image
        BufferedImage image = project.getImages().elementAt(index).getBufferedImage();

        int scale = 1;

        while (true) {

            if (image.getWidth() / scale > 320)
                scale++;
            else break;

        }

        int scaledWidth = image.getWidth() / scale;
        int scaledHeight = image.getHeight() / scale;


        data.setImage(trainingData.elementAt(index));
        int[][] overlay = new int[scaledWidth + scale][scaledHeight + scale];

        int overlayX;
        int overlayY = 0;

        for (int y = 0; y < image.getHeight(); y+=scale) {
            overlayX = 0;
            for (int x = 0; x < image.getWidth(); x+=scale) {
                try {
                    data.setX(x);
                    data.setY(y);
                    overlay[overlayX][overlayY] = (int) ind.execute(data);
                    overlayX++;
                } catch (ArrayIndexOutOfBoundsException e) {
                    // do nothing, oops.
                }
            }
            overlayY++;
        }

        BufferedImage out = new BufferedImage(scaledWidth, scaledHeight, BufferedImage.TYPE_INT_RGB);

        Graphics g = out.getGraphics();

        for (int y = 0; y < scaledHeight; y++) {
            for (int x = 0; x < scaledWidth; x++) {
                try {
                    g.setColor(colors[overlay[x][y]]);
                } catch (Exception e) {
                    //System.out.println("NO COLOUR FOR OVERLAY: " + overlay[x][y]);
                    g.setColor(Color.GRAY);
                }
                g.drawLine(x,y,x,y);
            }
        }

        return out;

    }


}