import org.jfree.chart.*;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.chart.plot.XYPlot;
import org.jfree.chart.renderer.xy.XYLineAndShapeRenderer;
import org.jfree.chart.title.TextTitle;
import org.jfree.data.xy.XYDataset;
import org.jfree.data.xy.XYSeries;
import org.jfree.data.xy.XYSeriesCollection;

import javax.imageio.ImageIO;
import javax.swing.*;
import javax.swing.border.Border;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.image.BufferedImage;
import java.io.*;
import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.util.ArrayList;

/**
 * Correlation Matrix Visualizer.
 * This program calculates correlations and visualizes them from given input file.
 *
 * @author Eerik Muuli
 */
public class CorrelationVisualizer {

    private static ArrayList<String> headers;
    private static int numberOfTotalRows;
    private static int matrixOrder;
    private static ArrayList<String[]> data;
    private static int screenshotCount = 0;
    private static int bigScreenshotCount = 0;

    public static void main(String[] args) {
        CorrelationVisualizer correlationVisualizer = new CorrelationVisualizer();
        correlationVisualizer.init();
    }

    /**
     * Core functionality is called out in this method.
     */
    private void init() {
        ArrayList<String> inputLines = readFile(); //kõik maatriksi read arraylistis.
        numberOfTotalRows = inputLines.size() - 1; //Maksimaalne ridade arv ühes veerus.
        headers = getHeaders(inputLines); //kõik pealkirjad ühes arraylistis.
        matrixOrder = headers.size(); //maatriksi külje pikkus (pealkirjade arv.
        data = processData(inputLines); //Ma vist ei kasuta seda muutujat tegelikult mitte kusagil.
        drawMatrix(); //Joonistab maatriksi.
    }

    /**
     * Draws the correlation matrix.
     * Also applies the listeners.
     */
    private static void drawMatrix() {
        JFrame frame = new JFrame("Correlation Matrix");
        frame.setDefaultCloseOperation(WindowConstants.EXIT_ON_CLOSE);
//        frame.setLayout(new GridLayout(1,1));

        final JPanel componentsPanel = new JPanel(); //Siin sees on maatriks.
        componentsPanel.setLayout(new GridLayout(matrixOrder, matrixOrder));
        componentsPanel.setSize(matrixOrder * 146, matrixOrder * 146);
        JScrollPane scrollpane = new JScrollPane(componentsPanel); //Scrollimise paneel.

        JMenuBar menuBar = initJMenuBar(componentsPanel);
        frame.setJMenuBar(menuBar);

        Dimension screen = Toolkit.getDefaultToolkit().getScreenSize();
        int contentWidth = componentsPanel.getWidth() + 25;
        int contentHeight = componentsPanel.getHeight() + (int) menuBar.getPreferredSize().getHeight() + 50;

        if (screen.width - 100 < contentWidth || screen.height - 100 < contentWidth) {
            frame.setExtendedState(Frame.MAXIMIZED_BOTH);
        } else {
            frame.setSize(contentWidth, contentHeight);
            frame.setLocation(screen.width/2-frame.getSize().width/2, screen.height/2-frame.getSize().height/2);
        }

        ChartPanel[] panels = new ChartPanel[matrixOrder * matrixOrder]; //Visuaalsed korrelatsioonid.

        for (int i = 0; i < matrixOrder; i++) {
            for (int j = 0; j < matrixOrder; j++) {
                if (i == j) { //Peadiagonaal.
                    createMainDiagonalCell(componentsPanel, i);
                }
                if (i < j) { //Arvulised korrelatsioonid.
                    createCorrelationCell(componentsPanel, i, j);
                }
                if (i > j) { //Diagrammi kujul korrelatsioonid.
                    final int labelNr = i * matrixOrder + j;
                    Border border = BorderFactory.createLineBorder(Color.BLACK, 1);

                    XYDataset dataset = createSampleDataset(i, j);
                    JFreeChart chart2 = ChartFactory.createXYLineChart(
                            String.valueOf(labelNr), //"Line Chart Demo 6",      // chart title
                            null, //"X",                      // x axis label
                            null, //"Y",                      // y axis label
                            dataset,                  // data
                            PlotOrientation.VERTICAL,
                            false,                     // include legend
                            false,                     // tooltips
                            false                     // urls
                    );
                    chart2.setTitle(new TextTitle(String.valueOf(labelNr), new Font(null, 0, 1))); //Peidab pealkirja ära, kuna see oli arvulisel kujul & kole.

                    XYPlot plot = (XYPlot) chart2.getPlot();
                    XYLineAndShapeRenderer renderer = new XYLineAndShapeRenderer();
                    renderer.setSeriesLinesVisible(0, false);
                    plot.setRenderer(renderer);

                    panels[labelNr] = new ChartPanel(chart2);
                    panels[labelNr].setBorder(border);

                    panels[labelNr].setPreferredSize(new Dimension(146, 146)); //Diagrammi mõõtmed, vaikimisi ka teiste lahtrid nüüd samad.

                    panels[labelNr].addChartMouseListener(new ChartMouseListener() { //Listener, kui klikitakse diagrammi peal.

                        /**
                         * Listener for mouse clicks on diagrams.
                         * Opens up diagrams in new window.
                         * @param e -
                         */
                        public void chartMouseClicked(ChartMouseEvent e) {
                            XYDataset dataset2 = createSampleDataset(labelNr / CorrelationVisualizer.matrixOrder, labelNr % CorrelationVisualizer.matrixOrder);
                            JFreeChart chart2 = ChartFactory.createXYLineChart(
                                    headers.get(labelNr % CorrelationVisualizer.matrixOrder) + " / " + headers.get(labelNr / CorrelationVisualizer.matrixOrder), //"Line Chart Demo 6",      // chart title
                                    headers.get(labelNr / CorrelationVisualizer.matrixOrder), //"X",                      // x axis label
                                    headers.get(labelNr % CorrelationVisualizer.matrixOrder), //"Y",                      // y axis label
                                    dataset2,                  // data
                                    PlotOrientation.VERTICAL,
                                    false,                     // include legend
                                    false,                     // tooltips
                                    false                     // urls
                            );

                            XYPlot plot = (XYPlot) chart2.getPlot();
                            XYLineAndShapeRenderer renderer = new XYLineAndShapeRenderer();
                            renderer.setSeriesLinesVisible(0, false);
                            plot.setRenderer(renderer);

                            ChartPanel uhuu = new ChartPanel(chart2);

                            final JFrame panell = new JFrame("Panel");
                            panell.setSize(700, 700);

                            Dimension dim = Toolkit.getDefaultToolkit().getScreenSize();
                            panell.setLocation(dim.width / 2 - panell.getSize().width / 2, dim.height / 2 - panell.getSize().height / 2);

                            final JPanel test = new JPanel(new BorderLayout());
                            test.add(uhuu);
                            panell.add(test);

                            JMenuBar menuBar = new JMenuBar();
                            JMenu menu = new JMenu("Menu");
                            menuBar.add(menu);

                            JMenuItem item2 = new JMenuItem("Screenshot");

                            item2.addActionListener(new ActionListener() {
                                @Override
                                /**
                                 * Takes screenshot of the diagram.
                                 */
                                public void actionPerformed(ActionEvent e) {

                                    try {
                                        screenshotCount += 1;
                                        BufferedImage image = new BufferedImage(test.getWidth(), test.getHeight(), BufferedImage.TYPE_INT_RGB);
                                        test.paint(image.getGraphics());
                                        File fail = new File("Panel" + screenshotCount + ".png");
                                        while (fail.exists()) {
                                            screenshotCount++;
                                            fail = new File("Panel" + screenshotCount + ".png");
                                        }
                                        ImageIO.write(image, "png", fail);
                                    } catch (Exception exception) {
                                        System.out.print("Exception while writing image");
                                    }
                                }
                            });
                            menu.add(item2);
                            panell.setJMenuBar(menuBar);
                            panell.setVisible(true);
                        }
                        public void chartMouseMoved(ChartMouseEvent e) {
                        }
                    });
                    componentsPanel.add(panels[labelNr]);
                }
            }
        }
        frame.add(scrollpane);
        frame.setVisible(true);
    }

    /**
     * Creates a new Jlabel with calculated correlation in it.
     * @param componentsPanel - the Panel in which the matrix is in.
     * @param i - the line of the matrix where the cell should be created.
     * @param j - the column of the matrix where the cell should be created.
     */
    private static void createCorrelationCell(JPanel componentsPanel, int i, int j) {
        NumberFormat formatter = new DecimalFormat("#0.00000"); //Et ei oleks lõputult pikad arvud.
        String x = formatter.format(calculateCorrelation(i, j)); //Arvutab korrelatsiooni.

        JLabel label = new JLabel(x);

        Border border = BorderFactory.createLineBorder(Color.BLACK, 1);
        label.setBorder(border);

        label.setHorizontalAlignment(JLabel.CENTER);

        componentsPanel.add(label);
    }

    /**
     * Creates a new JLabel with the attribute name in it.
     * @param componentsPanel - the Panel in which the matrix is in.
     * @param i - the row and the column in the matrix where the header cell should be created.
     */
    private static void createMainDiagonalCell(JPanel componentsPanel, int i) {
        JLabel label = new JLabel(headers.get(i));

        Border border = BorderFactory.createLineBorder(Color.BLACK, 1);
        label.setBorder(border); //Ääris lahtrite vahel.

        label.setHorizontalAlignment(JLabel.CENTER); //Tekst keskele.
        label.setFont(new Font(null, Font.BOLD, 15));

        componentsPanel.add(label);
    }

    /**
     * Creates JMenuBar with Exit and Screenshot options.
     * @param componentsPanel - the Panel in which the matrix is in.
     * @return - returns JMenuBar.
     */
    private static JMenuBar initJMenuBar(final JPanel componentsPanel) {
        JMenuBar menuBar = new JMenuBar();
        JMenu menu = new JMenu("Menu");
        menuBar.add(menu);

        JMenuItem item2 = new JMenuItem("Screenshot");
        item2.addActionListener(new ActionListener() {
            @Override
            /**
             * Takes screenshot of the Matrix.
             */
            public void actionPerformed(ActionEvent e) {
                try {
                    bigScreenshotCount++;
                    BufferedImage image = new BufferedImage(componentsPanel.getWidth(), componentsPanel.getHeight(), BufferedImage.TYPE_INT_RGB);
                    componentsPanel.paint(image.getGraphics());
                    File fail = new File("Matrix" + bigScreenshotCount+ ".png");
                    while (fail.exists()) {
                        bigScreenshotCount++;
                        fail = new File("Matrix" + bigScreenshotCount + ".png");
                    }
                    ImageIO.write(image, "png", fail);
                } catch (Exception exception) {
                    System.out.print("Exception while writing image");
                }
            }
        });
        menu.add(item2);

        JMenuItem item = new JMenuItem("Exit");
        item.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                System.exit(0);
            }
        });
        menu.add(item);
        return menuBar;
    }

    /**
     * Gets data from given rows and creates a dataset.
     * @param row1 - First row from which data must be taken.
     * @param row2 - Second row from which data must be taken.
     * @return - retuns a dataset.
     */
    private static XYDataset createSampleDataset(int row1, int row2) {
        final XYSeries series1 = new XYSeries("First");

        int minSize = 0;
        for (int i = 0; i < numberOfTotalRows; i++) {
            if ((data.get(i)[row1].equals("")) || (data.get(i)[row2].equals(""))) {
                minSize = i;
                break;
            }
        }

        for (int i = 0; i < minSize; i++) {
            String local1 = data.get(i)[row1].replace(",", ".");
            String local2 = data.get(i)[row2].replace(",", ".");
            series1.add(Double.parseDouble(local1), Double.parseDouble(local2));
        }

        final XYSeriesCollection dataset = new XYSeriesCollection();
        dataset.addSeries(series1);
        return dataset;
    }

    /**
     * Reads in the data from selected input file.
     * If the input file is not correct or cancel is pressed a sample dataset will be used.
     * @return - returns array of input lines.
     */
    private ArrayList<String> readFile() {

        JFileChooser fileChooser = new JFileChooser();
        fileChooser.setCurrentDirectory(new File(""));
        int chooseFile = fileChooser.showOpenDialog(null);

        BufferedReader reader = new BufferedReader(new InputStreamReader(
                this.getClass().getClassLoader().getResourceAsStream("test.csv")));

        if (chooseFile == JFileChooser.APPROVE_OPTION) {
            File selectedFile = fileChooser.getSelectedFile();
            try {
                reader = new BufferedReader(new FileReader(selectedFile));
            } catch (FileNotFoundException e) {
                e.printStackTrace();
            }
            if (!selectedFile.getName().substring(selectedFile.getName().length()- 4).equals(".csv")) {
                JOptionPane.showMessageDialog(null, "You didn't choose a CSV file, will use test.csv for now.");
                reader = new BufferedReader(new InputStreamReader(
                        this.getClass().getClassLoader().getResourceAsStream("test.csv")));
            }
        }

        ArrayList<String> fileData = new ArrayList<>();

        try (BufferedReader br = reader) {
            String line;
            while ((line = br.readLine()) != null) {
                fileData.add(line);
            }
        } catch (IOException e) {
            System.out.println(e.toString());
            System.exit(1);
        }
        return fileData;
    }

    /**
     * Gets headers from the given arraylist (first row).
     * @param input - arraylist with input file lines.
     * @return - returns an array with input file first row elements.
     */
    private static ArrayList<String> getHeaders(ArrayList<String> input) {
        ArrayList<String> mainDiagonal = new ArrayList<>();

        String[] parts = input.get(0).split(";");

        for (String part : parts) {
            mainDiagonal.add(part.toUpperCase()); //Siin on kõigi tulpade pealkirjad.
        }
        return mainDiagonal;
    }

    /**
     * Splits the lines in the given data.
     * @param input - arraylist with input lines.
     * @return - returns arraylist of String lists whitch consists of elements.
     */
    private static ArrayList<String[]> processData(ArrayList<String> input) {
        input.remove(0); //Eemaldame pealkirjad.

        ArrayList<String[]> processedData = new ArrayList<>();

        for (String in : input) {
            String[] parts = in.split(";");
            if (parts.length < matrixOrder) {
                String[] parts2 = new String[matrixOrder];
                for (int j = 0; j < matrixOrder; j++) {
                    parts2[j] = j < parts.length ? parts[j] : "";
                }
                parts = parts2;
            }
            processedData.add(parts);
        }
        return processedData;
    }

    /**
     * Calculates the correlation between the elements of 2 columns.
     * If one of the columns is shorter than the other, then
     * the calculations between pairs will stop at the shorter list last element.
     * @param row1 - First column that will be used in calculations.
     * @param row2 - Second column that will be used in calculations.
     * @return - returns the correlation between the 2 columns.
     */
    private static double calculateCorrelation(int row1, int row2) {
        ArrayList<Double> row1Values = new ArrayList<>();
        ArrayList<Double> row2Values = new ArrayList<>();
        ArrayList<Double> row1TimesRow2Values = new ArrayList<>();
        ArrayList<Double> row1TimesRow1Values = new ArrayList<>();
        ArrayList<Double> row2TimesRow2Values = new ArrayList<>();
        double sumRow1 = 0.0;
        double sumRow2 = 0.0;
        double sumRow1TimesRow2 = 0.0;
        double sumRow1TimesRow1 = 0.0;
        double sumRow2TimesRow2 = 0.0;
        double answer;

        int minSize;

        minSize = 0;
        for (int i = 0; i < numberOfTotalRows; i++) {
            if ((data.get(i)[row1].equals("")) || (data.get(i)[row2].equals(""))) {
                minSize = i;
                break;
            }
        }

        for (int i = 0; i < minSize; i++) {
            String local1 = data.get(i)[row1].replace(",", ".");
            String local2 = data.get(i)[row2].replace(",", ".");
            row1Values.add(Double.parseDouble(local1));
            row2Values.add(Double.parseDouble(local2));
        }

        for (int i = 0; i < row1Values.size(); i++) {
            row1TimesRow2Values.add(row1Values.get(i) * row2Values.get(i)); //x*y
            row1TimesRow1Values.add(row1Values.get(i) * row1Values.get(i)); //x*x
            row2TimesRow2Values.add(row2Values.get(i) * row2Values.get(i)); //y*y
            sumRow1 += row1Values.get(i);
            sumRow2 += row2Values.get(i);

        }

        for (int i = 0; i < row1TimesRow2Values.size(); i++) {
            sumRow1TimesRow2 += row1TimesRow2Values.get(i);
            sumRow1TimesRow1 += row1TimesRow1Values.get(i);
            sumRow2TimesRow2 += row2TimesRow2Values.get(i);
        }

        answer = ((row1Values.size() * sumRow1TimesRow2) - (sumRow1 * sumRow2)) /
                Math.sqrt(((row1Values.size() * sumRow1TimesRow1) - (sumRow1 * sumRow1))
                        * ((row1Values.size() * sumRow2TimesRow2) - (sumRow2 * sumRow2)));

        return answer;
    }

}
