package ac.essex.gp.interfaces;

import ac.essex.gp.params.GPParams;
import ac.essex.gp.individuals.Individual;
import ac.essex.gp.problems.*;
import ac.essex.gp.problems.picking.ClassifierPickingProblem2;
import ac.essex.gp.util.*;
import ac.essex.gp.Evolve;
import ac.essex.gp.training.segmentation.TrainingImage;

import javax.swing.*;
import java.awt.*;
import java.awt.datatransfer.Transferable;
import java.awt.datatransfer.StringSelection;
import java.awt.datatransfer.Clipboard;
import java.awt.event.*;
import java.awt.image.BufferedImage;
import java.util.Vector;
import java.text.DecimalFormat;

import org.jfree.chart.ChartFactory;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.category.DefaultCategoryDataset;

/**
 * Displays the progress of GP in real time in a graphical interface.
 *
 * @author Olly Oechsle, University of Essex, Date: 17-Jan-2007
 * @version 1.0
 */
public class GraphicalInterface extends GPInterface {

    protected EvolutionWindow window;
    protected ImageWindow imageWindow;
    protected ClassResultsWindow classWindow;
    protected long startTime;
    protected boolean displayImage = true;
    Vector<GenerationResult> results = new Vector<GenerationResult>(10);

    int counter = 0;
    int index = 0;
    int subgeneration = 0;

    Evolve e;
    Problem p;

    public GraphicalInterface() {
        try {
            UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName());
        } catch (Exception ex) {
            System.out.println("Unable to load native look and feel");
        }
    }

    public static final int NORMAL_GP = 1;
    public static final int SUB_GENERATIONAL_GP = 2;

    int graphMode = NORMAL_GP;

    public void writeToClipboard(String writeMe) {
        // get the system clipboard
        Clipboard systemClipboard = Toolkit.getDefaultToolkit().getSystemClipboard();

        // set the textual content on the clipboard
        Transferable transferableText = new StringSelection(writeMe);
        systemClipboard.setContents(transferableText, null);

    }

    public void dispose() {
        if (window != null) window.dispose();
        if (imageWindow != null) imageWindow.dispose();
        if (classWindow != null) classWindow.dispose();
    }

    public void onStartEvolution(Evolve e, Problem p) {
        this.e = e;
        this.p = p;
        final String title = p.getName();
        startTime = System.currentTimeMillis();
        if (p instanceof ClassifierPickingProblem2) {
            graphMode = SUB_GENERATIONAL_GP;
        }
        window = new EvolutionWindow();
        window.setTitle(title + " - Genetic Programming");
    }

    public void onStopped() {
        long time = (System.currentTimeMillis() - startTime) / 1000;
        JOptionPane.showMessageDialog(null, "Genetic Programming stopped.\n" + time + " seconds elapsed.");
    }

    public void onGenerationStart() {
        if (window != null) window.updateGenerationNumber(currentGeneration);
    }

    long lastTime = 0;

    public void onGoodIndividual(Individual ind) {
        if (window != null) {
            window.displayIndividual(ind);
            subgeneration++;
            if (bestIndividual != null)  {
                results.add(new GenerationResult(currentGeneration, bestIndividual, System.currentTimeMillis() - startTime));
            }
        }
    }

    public void incrementIndividualEvaluations() {
        super.incrementIndividualEvaluations();
        if (System.currentTimeMillis() - lastTime > 1000) {
            if (window != null) {
                window.individuals.setText("Individuals: " + getTotalIndividualsEvaluated());
                window.updateTime();
            }
            lastTime = System.currentTimeMillis();
        }
    }

    double previousFitness = -1;

    boolean alreadyWorking = false;

    public synchronized void onGenerationEnd() {

        if (bestIndividual != null)  {
            results.add(new GenerationResult(currentGeneration, bestIndividual, System.currentTimeMillis() - startTime));
        }

        if (window != null) {
            window.individuals.setText("Individuals: " + getTotalIndividualsEvaluated());
            window.displayIndividual(bestIndividual);

            final int generationNumber = currentGeneration;

            if (p instanceof JasmineSegmentationProblem) {

                if (displayImage && !alreadyWorking) {

                    SwingUtilities.invokeLater(new Thread() {
                        public void run() {
                            alreadyWorking = true;
                            if (bestIndividual.getKozaFitness() != previousFitness || imageWindow.image.getIcon() == null) {
                                if (imageWindow == null) imageWindow = new ImageWindow();
                                imageWindow.lblMessage.setText("Drawing...");
                                previousFitness = bestIndividual.getKozaFitness();
                                Object o = p.describe(bestIndividual, new DataStack(), index);
                                if (o != null && o instanceof BufferedImage) {
                                    imageWindow.setImage((BufferedImage) o);
                                }
                                imageWindow.lblMessage.setText("Result of generation: " + generationNumber);
                            }
                            alreadyWorking = false;
                        }
                    });
                }

            }


            if (p instanceof ShapeClassificationProblem) {

                ac.essex.gp.util.ClassResults results = (ac.essex.gp.util.ClassResults) p.describe(bestIndividual, new DataStack(), index);

                if (classWindow == null)
                    classWindow = new ClassResultsWindow(results);
                else
                    classWindow.update(results);

            }
        }
    }

    public void onEndEvolution(GPParams params) {
        if (isIdeal) {
            JOptionPane.showMessageDialog(window, "Found ideal individual", "Complete!", JOptionPane.INFORMATION_MESSAGE);
        } else {
            JOptionPane.showMessageDialog(window, "Reached " + currentGeneration + " generations.", "Complete!", JOptionPane.INFORMATION_MESSAGE);
        }
    }

    class ImageMenu extends JMenuItem {

        int i;

        public ImageMenu(ac.essex.gp.training.segmentation.TrainingImage image, int i) {

            super("Image " + i);

            this.i = i;

            addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    previousFitness = -1;
                    index = ImageMenu.this.i;
                }
            });

        }

    }

    class ClassResultsWindow extends JDialog {

        public static final int TOTAL = 0;
        public static final int CORRECT = 1;
        public static final int PERCENT = 2;

        JLabel[][] classResults;

        public ClassResultsWindow(ac.essex.gp.util.ClassResults results) {

            super(window);

            setTitle("Live Results");

            Container c = getContentPane();

            classResults = new JLabel[results.classes.size()][3];

            c.setLayout(new GridLayout(results.classes.size() + 1, 4));

            c.add(new JLabel("Class"));
            c.add(new JLabel("Total"));
            c.add(new JLabel("Correct"));
            c.add(new JLabel("%"));

            for (int row = 0; row < results.classes.size(); row++) {
                ac.essex.gp.util.ClassResult r = results.classes.elementAt(row);
                c.add(new JLabel(r.name));

                classResults[row][TOTAL] = new JLabel("?");
                c.add(classResults[row][TOTAL]);

                classResults[row][CORRECT] = new JLabel("?");
                c.add(classResults[row][CORRECT]);

                classResults[row][PERCENT] = new JLabel("?");
                c.add(classResults[row][PERCENT]);

            }

            setSize(300, (20 * results.classes.size()) + 40);
            setVisible(true);

            update(results);

        }

        public void update(ac.essex.gp.util.ClassResults results) {

            DecimalFormat f = new DecimalFormat("0.0");

            for (int row = 0; row < results.classes.size(); row++) {
                ac.essex.gp.util.ClassResult c = results.classes.elementAt(row);

                classResults[row][TOTAL].setText(String.valueOf(c.total));

                classResults[row][CORRECT].setText(String.valueOf(c.correct));

                classResults[row][PERCENT].setText(f.format((c.correct / (double) c.total) * 100));

            }

        }

    }

    class ImageWindow extends JDialog {

        JLabel image;
        JLabel lblMessage;

        public ImageWindow() {
            super(window);
            setTitle("Result Viewer");
            image = new JLabel();
            lblMessage = new JLabel();

            JPanel status = new JPanel(new FlowLayout(FlowLayout.LEFT));
            getContentPane().add(image, BorderLayout.CENTER);
            getContentPane().add(lblMessage, BorderLayout.SOUTH);
            setLocation((int) window.getLocation().getX(), 50);
            setVisible(true);

            addWindowListener(new WindowAdapter() {
                public void windowClosing(WindowEvent e) {
                    onClose();
                }
            });

        }

        public void setImage(BufferedImage img) {
            setSize(img.getWidth(), img.getHeight() + 50);
            image.setIcon(new ImageIcon(img));
        }

        public void onClose() {
            displayImage = false;
            window.mnuDisplayImage.setSelected(false);
            dispose();
        }

    }

    class EvolutionWindow extends JFrame {

        JLabel generationNumber;
        JLabel fitness;
        JLabel hits;
        JLabel time;
        JLabel individuals;

        JMenuItem mnuExit, mnuExportStats;
        JCheckBoxMenuItem mnuDisplayImage;

        JLabel chart;
        JTextArea displayIndividual;
        DefaultCategoryDataset series;

        JButton smack;

        public EvolutionWindow() {

            series = new DefaultCategoryDataset();

            Container c = getContentPane();
            generationNumber = new JLabel("Initialising...");
            fitness = new JLabel();
            hits = new JLabel();
            time = new JLabel();
            individuals = new JLabel();

            chart = new JLabel();
            c.add(chart, BorderLayout.CENTER);

            displayIndividual = new JTextArea("Evaluating: Generation 0");
            JScrollPane pane = new JScrollPane(displayIndividual);
            pane.setPreferredSize(new Dimension(350, -1));
            c.add(pane, BorderLayout.EAST);

            // MENUS

            mnuDisplayImage = new JCheckBoxMenuItem("Preview Image");
            mnuDisplayImage.setAccelerator(KeyStroke.getKeyStroke(KeyEvent.VK_F1, 0));
            mnuDisplayImage.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    if (!mnuDisplayImage.isSelected()) {
                        imageWindow.onClose();
                    } else {
                        previousFitness = -1;
                        imageWindow = new ImageWindow();
                    }
                    displayImage = mnuDisplayImage.isSelected();
                }
            });
            mnuDisplayImage.setSelected(displayImage);

            mnuExportStats = new JMenuItem("Export Statistics");
            mnuExportStats.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    new StatisticsDisplay(results);
                }
            });

            JMenuItem mnuStop = new JMenuItem("Stop");
            mnuStop.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    GraphicalInterface.this.e.stopFlag = true;
                    ((JMenuItem) e.getSource()).setEnabled(false);
                }
            });

            mnuExit = new JMenuItem("Exit");
            mnuExit.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    onExit();
                }
            });

            JMenuItem mnuCopyToClipboard = new JMenuItem("Copy code to clipboard");
            mnuCopyToClipboard.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    writeToClipboard(displayIndividual.getText());
                }
            });

            JMenu edit = new JMenu("Edit");
            edit.add(mnuCopyToClipboard);


            smack = new JButton("Smack");
            smack.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    GraphicalInterface.this.e.smack = true;
                }
            });

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


            JMenu view = new JMenu("View");
            view.add(mnuDisplayImage);

            JMenuBar bar = new JMenuBar();
            bar.add(file);
            bar.add(edit);
            bar.add(view);

            if (p instanceof JasmineSegmentationProblem) {
                Vector<ac.essex.gp.training.segmentation.TrainingImage> trainingData = ((JasmineSegmentationProblem) p).trainingData;
                JMenu preview = new JMenu("Preview");
                for (int i = 0; i < trainingData.size(); i++) {
                    TrainingImage trainingImage = trainingData.elementAt(i);
                    preview.add(new ImageMenu(trainingImage, i));
                }
                bar.add(preview);
            }

            setJMenuBar(bar);

            // END MENUS

            JPanel p = new JPanel(new GridLayout(1, 6));


            p.add(generationNumber);
            p.add(individuals);
            p.add(fitness);
            p.add(hits);
            p.add(time);
            p.add(smack);


            c.add(p, BorderLayout.SOUTH);
            setTitle("Genetic Programming");

            addWindowListener(new WindowAdapter() {
                public void windowClosing(WindowEvent e) {
                    onExit();
                }
            });

            setSize(700, 300);
            setLocation((1280 - 700) / 2, (1024 - 300) / 2);
            setVisible(true);

        }

        public void onExit() {
            GraphicalInterface.this.e.stopFlag = true;
            if (imageWindow != null) imageWindow.dispose();
            dispose();
        }

        public void updateGenerationNumber(int num) {
            generationNumber.setText("Generation: " + num);
        }

        public void displayIndividual(Individual ind) {
            if (ind != null) {
                if (graphMode == NORMAL_GP) {
                    String java = ind.toJava();
                    if (!java.equals(displayIndividual.getText())) {
                        displayIndividual.setText(java);
                        displayIndividual.setSelectionStart(java.length() - 10);
                        displayIndividual.setSelectionEnd(java.length() - 1 );
                    }
                } else {
                    ClassifierPickingProblem2 cpp = (ClassifierPickingProblem2) p;
                    displayIndividual.setText(cpp.buffer.toString());
                }
                fitness.setText("Errors: " + ind.getKozaFitness());
                hits.setText(" Hits: " + ind.getHits());
                updateChart(ind.getKozaFitness(), ind.getHits(), currentGeneration);
            }
        }

        public void updateTime() {
            long elapsed = (System.currentTimeMillis() - startTime) / 1000;
            time.setText("Time: " + elapsed + " secs.");
        }

        public void updateChart(double fitness, int hits, int generation) {
            if (graphMode == NORMAL_GP) {
                series.addValue(fitness, "row1", String.valueOf(generation));
            } else {
                series.addValue(hits, "row1", String.valueOf(subgeneration));
            }
            JFreeChart myChart = ChartFactory.createLineChart(null, graphMode == NORMAL_GP ? "Generations" : "Sub Generations", graphMode == NORMAL_GP ? "Errors" : "Hits", series, PlotOrientation.VERTICAL, false, false, false);
            BufferedImage image = myChart.createBufferedImage(chart.getWidth(), chart.getHeight());
            chart.setIcon(new ImageIcon(image));
        }

    }

    public void fatal(String message) {
        JOptionPane.showMessageDialog(window, message, "Error", JOptionPane.ERROR_MESSAGE);
    }

    public void message(String message) {
        JOptionPane.showMessageDialog(window, message, "Message", JOptionPane.INFORMATION_MESSAGE);
    }

    class StatisticsDisplay extends JDialog {

        JTextArea area;

        public StatisticsDisplay(Vector<GenerationResult> results) {

            area = new JTextArea();
            JScrollPane pane = new JScrollPane(area);

            StringBuffer text = new StringBuffer(GenerationResult.getCSVHeader());
            text.append("\n");

            int lastResultHits = -1;

            for (int i = 0; i < results.size(); i++) {
                GenerationResult generationResult = results.elementAt(i);
                if (generationResult.individual.getHits() != lastResultHits) {
                    lastResultHits = generationResult.individual.getHits();
                    text.append(generationResult.toCSV());
                    text.append("\n");
                }
            }
                             
            area.setText(text.toString());

            getContentPane().add(pane, BorderLayout.CENTER);

            JPanel buttons = new JPanel(new FlowLayout(FlowLayout.RIGHT));

            JButton ok = new JButton("OK");

            ok.addActionListener(new ActionListener() {
                public void actionPerformed(ActionEvent e) {
                    StatisticsDisplay.this.dispose();
                }
            });

            addWindowListener(new WindowAdapter() {
                public void windowClosing(WindowEvent e) {
                    StatisticsDisplay.this.dispose();
                }
            });

            buttons.add(ok);

            getContentPane().add(buttons, BorderLayout.SOUTH);

            setTitle("GP Summary");
            setSize(300, 340);
            setLocation(200, 200);
            setVisible(true);


        }

    }

}