package ac.essex.ooechs.facedetection.util.evaluation;

import ac.essex.ooechs.imaging.commons.PixelLoader;
import ac.essex.ooechs.imaging.commons.ImageWindow;
import ac.essex.ooechs.imaging.commons.Pixel;
import ac.essex.ooechs.imaging.commons.apps.webcam.WebcamGrabber;
import ac.essex.ooechs.imaging.commons.segmentation.Segmenter;
import ac.essex.ooechs.imaging.commons.util.panels.ImagePanel;
import ac.essex.ooechs.imaging.shapes.Grouper;
import ac.essex.ooechs.imaging.shapes.SegmentedShape;
import ac.essex.ooechs.imaging.shapes.ExtraShapeData;
import ac.essex.ooechs.facedetection.singlestage.nodes.FeatureUtils;

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 = 30;
    public static int windowHeight = 30;

    protected JLabel result, position, status;

    protected File imageDirectory;

    protected boolean smiling = false;

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

    protected JTextField saveTo;
    protected JTextField badName;

    protected JMenuItem mnuExit, mnuRun, runAll;

    protected PixelLoader image;

    protected ImagePanel ip;

    protected Segmenter segmenter;
    protected ObjectDetector objectDetector;
    protected ObjectDetector smileDetector = new SmileDetector();
    //protected EyeDetector eyeDetector = new EyeDetector();

    protected JCheckBox saveEnabled;

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

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

    protected JButton btnCapture;
    protected JTextField time;

    public WebcamGrabber webcam;

    boolean stop;

    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/Data/me/originals"), new SkinSegmenter(), new MouthDetectorBoosted());

    }

/*    public static ObjectDetector getDetector(final AdaBoostSolution s) {
    return new ObjectDetector() {
        public double calculate(int x, int y, HaarlikeFeatures image) {
            try {
                return s.classify(new AdaBoostSample(image), null);
            } catch (ArrayIndexOutOfBoundsException e) {
                return ObjectDetector.NOT_OBJECT;
            }
        }

    };
}*/

    public WindowSearchGUI(File imageDirectory, Segmenter segmenter, ObjectDetector solution) {
        this(imageDirectory, segmenter, solution, "GP Smile Detection v1.0");
    }

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

        super(title);

        this.segmenter = segmenter;

        this.objectDetector = solution;

        if (imageDirectory.exists()) {
            this.imageDirectory = imageDirectory;
        } else {
            this.imageDirectory = new File(System.getProperty("user.home"));
        }

        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/Data/me/");

        badName = new JTextField("bad");

        buttons.add(saveTo);
        buttons.add(badName);
        buttons.add(new JLabel(".bmp"));
        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);*/

        saveEnabled = new JCheckBox();
        buttons.add(saveEnabled);

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

        try {
            webcam = new WebcamGrabber();

            btnCapture = new JButton("Test on Webcam");
            btnCapture.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    if (stop) {
                        stop = false;
                        Thread t = new Thread() {
                            public void run() {
                                while (!stop) {
                                    try {
                                        int t = Integer.parseInt(time.getText());
                                        sleep(t);
                                    } catch (Exception e) {
                                        // do nothing
                                    }
                                    windows = null;
                                    image = webcam.grab();
                                    showingMask = true;
                                    runFast();
                                    ip.setImage(image);
                                }
                            }
                        };
                        t.start();
                        btnCapture.setText("Stop");
                    } else {
                        stop = true;
                        btnCapture.setText("Test on Webcam");
                    }
                }
            });

            buttons.add(btnCapture);
            time = new JTextField("400");
            buttons.add(new JLabel("Interval: "));
            buttons.add(time);

        } catch (Exception e) {

        } catch (Error f) {

        }

        // 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 {
            System.out.println("saveWindows()");
            if (windows != null) {
                for (int i = 0; i < windows.size(); i++) {
                    ImageWindow pixel = windows.elementAt(i);
                    // only take 0.10%
                    //if (Math.random() > 0.5) continue;
                    File saveDir = new File(saveTo.getText());
                    if (!saveDir.exists()) {
                        saveDir.mkdirs();
                    }
                    if (saveDir.exists()) {
                        File f = new File(saveDir, badName.getText() + counter + ".bmp");
                        PixelLoader subimage = image.getSubImage(new ImageWindow(pixel.left, pixel.top, windowWidth, windowHeight));
                        subimage.saveAs(f);
                        counter++;
                    } else {
                        alert(saveTo.getText() + " is not a directory");
                    }
                }
                System.out.println("Counter: " + counter);
            }
        } catch (Exception e) {
            alert(e.toString());
        }
    }

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

    private void load(boolean show) {

        windows = null;

        results = null;

        try {

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

            // stop showing mask
            showingMask = false;

            // 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;

    int[][] results;

    public void runFast() {

        if (windows == null) {
            windows = new Vector<ImageWindow>(10);
        } else {
            windows.clear();
        }

        long start = System.currentTimeMillis();
        results = segmenter.segment(image);
        long time = System.currentTimeMillis() - start;

        //status.setText("Found " + windows.size() + " objects in " + time + "ms.");
        System.out.println("Segmented in " + time + " ms");

        Grouper gr = new Grouper();

        Vector<SegmentedShape> shapes = gr.getShapes(results);

        Graphics g = image.getBufferedImage().getGraphics();

        start = System.currentTimeMillis();

        windows = new Vector<ImageWindow>(10);

        int shapeCount = 0;

        for (int i = 0; i < shapes.size(); i++) {
            SegmentedShape segmentedShape = shapes.elementAt(i);

            if (segmentedShape.getMass() > 1000) {

                g.setColor(Color.BLUE);
                //g.drawString(segmentedShape.getMass() + "px", segmentedShape.minX, segmentedShape.minY);

                shapeCount++;
                int width = segmentedShape.maxX - segmentedShape.minX;
                int height = segmentedShape.maxY - segmentedShape.minY;

                windows.add(new ImageWindow(segmentedShape.minX, segmentedShape.minY, width, height));
/*
                if (saveEnabled.isSelected()) {

                    PixelLoader subimage = image.getSubImage(new ImageWindow(segmentedShape.minX, segmentedShape.minY, width, height));

                    File saveDir = new File(saveTo.getText());
                    if (!saveDir.exists()) {
                        saveDir.mkdirs();
                    }
                    if (saveDir.exists()) {
                        File f = new File(saveDir, badName.getText() + counter + ".bmp");
                        try {
                            subimage.saveAs(f);
                            counter++;
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    } else {
                        alert(saveTo.getText() + " is not a directory");
                    }

                }*/

                // find mouths within the window
                int[][] output = new int[width][height];

                for (int y = 0; y < height; y++) {
                    for (int x = 0; x < width; x++) {

                        try {
                            
                            if (eval(segmentedShape.minX + x, segmentedShape.minY + y) > 0) {
                                //windows.add(new ImageWindow(x + segmentedShape.minX, y + segmentedShape.minY, 1, 1));
                                output[x][y] = 1;
                            }

                        } catch (Exception e) {
                            // do nothing
                        }

                    }

                }

                Vector<SegmentedShape> mouthshapes = new Grouper().getShapes(output);

                SegmentedShape biggest = null;
                for (int j = 0; j < mouthshapes.size(); j++) {
                    SegmentedShape shape = mouthshapes.elementAt(j);
                    if (biggest == null || shape.getMass() > biggest.getMass()) {
                        biggest = shape;
                        //windows.add(new ImageWindow(segmentedShape.minX + biggest.minX, segmentedShape.minY + biggest.minY, windowWidth, windowHeight)) ;
                    }
                }

                FeatureUtils.reset();

                if (biggest != null) {
                    try {
                        ExtraShapeData esd = new ExtraShapeData(biggest);
                        int wx = segmentedShape.minX + biggest.minX + esd.getCentreOfGravityX();
                        int wy = segmentedShape.minY + biggest.minY + esd.getCentreOfGravityY();

                        int smile = (int) smileDetector.calculate(wx, wy, windowWidth, windowHeight, image);

                        if (smile == 1) {
                            smiling = true;
                            System.out.println("Smiling");
                        } else {
                            smiling = false;
                            System.out.println("Not smiling");
                        }

                        windows.add(new ImageWindow(wx, wy, windowWidth, windowHeight));


                    } catch (Exception e) {

                    }
                }

            }
        }

        time = System.currentTimeMillis() - start;

        //status.setText("Found " + windows.size() + " objects in " + time + "ms.");
        System.out.println("Classified in " + time + " ms");

        //}

        //}

/*        if (saveEnabled.isSelected()) {
            File saveDir = new File(saveTo.getText());
            if (!saveDir.exists()) {
                saveDir.mkdirs();
            }
            if (saveDir.exists()) {
                File f = new File(saveDir, "frame" + counter + ".bmp");
                try {
                    image.saveAs(f);
                    counter++;
                } catch (Exception e) {
                    e.printStackTrace();
                }
            } else {
                alert(saveTo.getText() + " is not a directory");
            }
        }*/

    }

    public void run() {

        runFast();


    }

    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++;
        windowHeight++;
        windows = null;
        ip.repaint();
    }

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

    public int eval(int x, int y) {
        return (int) objectDetector.calculate(x, y, windowWidth, windowHeight, image);
    }

    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 == ObjectDetector.OBJECT) {
                            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) {

            if (results != null) {
                g.setColor(new Color(255,255,255,128));
                for (int y = 0; y < image.getHeight(); y++)
                for (int x = 0; x < image.getWidth(); x++) {
                    if (results[x][y] > 0)  {
                        image.setRGB(x,y,Color.WHITE.getRGB());
                    } else {
                        image.setRGB(x,y,Color.BLACK.getRGB());
                    }
                }
            }

            super.paintComponent(g);

/*            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.BLUE);
                    }

                    drawRect(g, window.left, window.top, window.width, window.height);
                }
            }

            drawRect(g, x, y, windowWidth, windowHeight);

            if (smiling == true) {
                g.setColor(Color.GRAY);
                g.fillRect(0, image.getHeight() - 40, image.getWidth(), 40);
                g.setColor(Color.WHITE);
                g.drawString("SMILING", image.getWidth() / 2, image.getHeight() - 20);
            }*/

        }

    }
}