package ac.essex.gp.ooechs.detection.util;

import ac.essex.ooechs.imaging.commons.PixelLoader;
import ac.essex.ooechs.imaging.commons.HaarRegions;
import ac.essex.ooechs.imaging.commons.ImageWindow;
import ac.essex.ooechs.imaging.commons.evolved.ObjectDetector;
import ac.essex.ooechs.imaging.commons.util.panels.ImagePanel;
import ac.essex.ooechs.adaboost.AdaBoostSolution;
import ac.essex.ooechs.adaboost.AdaBoostSample;
import ac.essex.gp.ooechs.detection.gp.DetectionProblem;

import javax.swing.*;
import java.awt.event.*;
import java.awt.*;
import java.io.File;
import java.io.FilenameFilter;
import java.util.Vector;

/**
 * Allows the user to evaluate subjectively the performance of a detection
 * algorithm and facilitates the production of better training data to
 * improve algorithms further.
 *
 * @author Olly Oechsle, University of Essex, Date: 19-Sep-2006
 * @version 1.1 28th April 2008
 */
public class WindowSearchGUI extends JFrame implements ActionListener {

    public static int windowWidth = DetectionProblem.WINDOW_WIDTH;
    public static int windowHeight = DetectionProblem.WINDOW_HEIGHT;

    public static int SEGMENTSX = DetectionProblem.WINDOWBLOCKSX;
    public static int SEGMENTSY = DetectionProblem.WINDOWBLOCKSY;

    protected JLabel result, position, status;

    protected File imageDirectory;

    protected JButton next, back, save, plus, minus;

    protected JTextField saveTo;

    protected JMenuItem mnuExit, mnuRun, runAll;

    protected PixelLoader image;

    protected ImagePanel ip;
    protected HaarRegions haar;
    protected ObjectDetector solution;

    // list of all the files in the directory
    protected File[] images;

    // index of our position through the array of files
    protected int index = 0;



    public static void main(String[] args) throws Exception {

        AdaBoostSolution s = AdaBoostSolution.load(new File("/home/ooechs/Desktop/birds/toucan/eyes.solution"));

        new WindowSearchGUI(new File("/home/ooechs/Desktop/birds/toucan"), getDetector(s));

    }

    public static ObjectDetector getDetector(final AdaBoostSolution s) {
        return new ObjectDetector() {
            public double calculate(HaarRegions haarRegions) {
                try {
                    return s.classify(new AdaBoostSample(haarRegions), null);
                } catch (ArrayIndexOutOfBoundsException e) {
                    return DetectionProblem.NOT_OBJECT;
                }
            }

        };
    }

    public WindowSearchGUI(File imageDirectory, ObjectDetector solution) {
        this(imageDirectory, solution, "Evolved Object Detection using Genetic Programming");
    }

    public WindowSearchGUI(File imageDirectory, ObjectDetector solution, String title) {

        super(title);

        this.solution = solution;

        this.imageDirectory = imageDirectory;

        loadImages();

        // SWING STUFF BELOW ---
        Container c = getContentPane();

        // MENU

        mnuExit = new JMenuItem("Exit");
        mnuExit.addActionListener(this);

        mnuRun = new JMenuItem("Run");
        mnuRun.addActionListener(this);

        runAll = new JMenuItem("Run All");
        runAll.addActionListener(this);

        JMenu file = new JMenu("File");
        file.add(mnuExit);

        JMenu tools = new JMenu("Tools");

        tools.add(mnuRun);
        tools.add(runAll);

        JMenuBar bar = new JMenuBar();

        bar.add(file);
        bar.add(tools);

        setJMenuBar(bar);

        // UPPER PART

        JToolBar buttons = new JToolBar();

        save = new JButton("Save");
        save.addActionListener(this);

        saveTo = new JTextField("/home/ooechs/Desktop/birds/mandarin/false");

        buttons.add(saveTo);
        buttons.add(save);

        buttons.addSeparator();

        next = new JButton("Next");
        back = new JButton("Prev");

        next.addActionListener(this);
        back.addActionListener(this);

        buttons.addSeparator();

        plus = new JButton("+");
        minus = new JButton("-");

        plus.addActionListener(this);
        minus.addActionListener(this);

        buttons.add(back);
        buttons.add(next);

        buttons.add(minus);
        buttons.add(plus);

        c.add(buttons, BorderLayout.NORTH);

        // CENTER PART
        ip = new WindowPanel();

        // load the first image in the directory
        load(true);

        c.add(new JScrollPane(ip), BorderLayout.CENTER);

        // BOTTOM PART
        status = new JLabel("");
        position = new JLabel("");
        result = new JLabel("");

        JPanel statusPanel = new JPanel(new GridLayout(1,3));

        statusPanel.add(status);
        statusPanel.add(position);
        statusPanel.add(result);

        c.add(statusPanel, BorderLayout.SOUTH);

        // EXIT JAVA WHEN WINDOW CLOSES
        addWindowListener(new WindowAdapter() {
            public void windowClosing(WindowEvent e) {
                exit();
            }
        });

        // SHOW WINDOW
        setSize(800, 600);
        setVisible(true);

    }

    public void exit() {
        System.exit(0);
    }

    public void actionPerformed(ActionEvent e) {

        if (e.getSource() == next) next(true);
        if (e.getSource() == back) back(true);
        if (e.getSource() == plus) increaseWindowSize();
        if (e.getSource() == minus) decreaseWindowSize();
        if (e.getSource() == mnuRun) {
            run();
            ip.repaint();
        }
        if (e.getSource() == mnuExit) exit();
        if (e.getSource() == save) {
            saveWindows();
        }

    }

    int counter = 0;

    private void saveWindows() {
        try {
            if (windows != null) {
                for (int i = 0; i < windows.size(); i++) {
                    ImageWindow pixel = windows.elementAt(i);
                    // only take 10%
                    //home//if (Math.random() > 0.1) continue;
                    File saveDir = new File(saveTo.getText());
                    if (!saveDir.exists()) {
                        saveDir.mkdirs();
                    }
                    if (saveDir.exists()) {
                        File f = new File(saveDir, "false" + counter + ".bmp");
                        PixelLoader subimage = haar.getPixelLoader().getSubImage(new ImageWindow(pixel.left, pixel.top, windowWidth, windowHeight));
                        subimage.saveAs(f);
                        counter++;
                    } else {
                        alert(saveTo.getText() + " is not a directory");
                    }
                }
            }
        } catch (Exception e) {
            alert(e.toString());
        }
    }

    public void alert(String message) {
        JOptionPane.showMessageDialog(this, message);
    }

    private void load(boolean show) {

        windows = null;

        try {

            // load image
            image = new PixelLoader(images[index]);

            // stop showing mask
            showingMask = false;

            // create the haar image
            haar = new HaarRegions(image);

            // show the image
            if (show) ip.setImage(image);

        } catch (Exception e) {
            JOptionPane.showMessageDialog(this, "Could not load image: " + e.getMessage());
            e.printStackTrace();
        }
    }

    protected Vector<ImageWindow> windows = null;

    public void run() {

        long start = System.currentTimeMillis();

        windows = new Vector<ImageWindow>(10);

        for (int y = 0; y < (image.getHeight() - haar.getWindowHeight()); y += 2) {
            for (int x = 0; x < (image.getWidth() - haar.getWindowWidth()); x += 2) {

                int result = eval(x, y);

                if (result == DetectionProblem.OBJECT1) {
                    windows.add(new ImageWindow(x, y, haar.getWindowWidth(), haar.getWindowHeight()));
                }

            }

        }

        long time = System.currentTimeMillis() - start;

        status.setText("Found " + windows.size() + " objects in " + time + "ms.");

    }

    public void loadImages() {

        images = imageDirectory.listFiles(new FilenameFilter() {
            public boolean accept(File dir, String name) {
                return name.endsWith(".bmp") || name.endsWith(".gif") || name.endsWith(".jpg");
            }
        });

        index = 0;
    }

    public boolean next(boolean show) {

        if (index < (images.length - 1)) {
            index++;
            load(show);
            return true;
        } else {
            return false;
        }

    }

    public boolean back(boolean show) {

        if (index > 1) {
            index--;
            load(show);
            return true;
        } else {
            return false;
        }

    }

    public void increaseWindowSize() {
        windowWidth += SEGMENTSX;
        windowHeight += SEGMENTSY;
        windows = null;
        ip.repaint();
    }

    public void decreaseWindowSize() {
        windowWidth -= SEGMENTSX;
        windowHeight -= SEGMENTSY;
        windows = null;
        ip.repaint();
    }

    public int eval(int x, int y) {

        haar.setWindowPosition(x, y, windowWidth, windowHeight, SEGMENTSX, SEGMENTSY);

        return (int) solution.calculate(haar);

    }

    boolean showingMask = true;

    private void toggle() {
        showingMask = !showingMask;
        if (showingMask) {
            if (windows == null) {
                run();
            }
        }
        ip.repaint();
    }

    class WindowPanel extends ImagePanel {

        protected int x, y;

        public WindowPanel() {

            addMouseMotionListener(new MouseMotionAdapter() {
                public void mouseMoved(MouseEvent e) {
                    x = getX(e);
                    y = getY(e);

                    if (showingMask && windows != null) {

                        for (int i = 0; i < windows.size(); i++) {
                            ImageWindow window = windows.elementAt(i);
                            if (contains(window, x, y)) {
                                window.colour = 1;
                            } else {
                                window.colour = 0;
                            }
                        }

                    }

                    position.setText("x=" + x + ", y=" + y);
                    repaint();

                    try {
                        int decision = eval(x, y);
                        if (decision == DetectionProblem.OBJECT1) {
                            result.setText(" FOUND OBJECT");
                        } else {
                            result.setText("");
                        }
                    } catch (RuntimeException e1) {
                        result.setText("Out of bounds");
                    }
                }
            });

            addMouseListener(new MouseAdapter() {
                public void mousePressed(MouseEvent e) {
                    if (e.getButton() == MouseEvent.BUTTON1) {
                        toggle();
                    } else {
                        if (windows != null) {
                            // delete the red windows
                            Vector<ImageWindow> toRemove = new Vector<ImageWindow>();
                            for (int i = 0; i < windows.size(); i++) {
                                ImageWindow pixel = windows.elementAt(i);
                                if (pixel.colour == 1) {
                                    toRemove.add(pixel);
                                }
                            }
                            windows.removeAll(toRemove);
                            repaint();
                        }
                    }
                }
            });

            addMouseWheelListener(new MouseWheelListener() {
                public void mouseWheelMoved(MouseWheelEvent e) {
                    if (e.getUnitsToScroll() > 0) {
                        zoomOut();
                    } else {
                        zoomIn();
                    }
                    repaint();
                }
            });

        }

        public boolean contains(ImageWindow window, int x, int y) {
            return (x >= window.left && x <= (window.left + windowWidth) && y >= window.top && y <= (window.top + windowHeight));
        }

        public void paintComponent(Graphics g) {
            super.paintComponent(g);

            if (!showingMask) {
                drawRect(g, x, y, windowWidth, windowHeight);
            }

            if (showingMask && windows != null) {
                for (int i = 0; i < windows.size(); i++) {
                    ImageWindow window = windows.elementAt(i);
                    if (window.colour == 1) {
                        g.setColor(Color.RED);
                    } else {
                        g.setColor(Color.WHITE);
                    }
                    drawRect(g, window.left, window.top, window.width, window.height);
                }
            }
        }

    }
}