package beastfx.app.inputeditor;


import java.awt.BorderLayout;
import java.awt.Color;
import java.awt.Dimension;
import java.awt.Font;
import java.io.File;
import java.util.HashMap;
import java.util.Map;

import javax.swing.Box;
import javax.swing.JCheckBox;
import javax.swing.JPanel;
import javax.swing.JScrollPane;
import javax.swing.JTable;
import javax.swing.JViewport;
import javax.swing.ListSelectionModel;
import javax.swing.ScrollPaneConstants;
import javax.swing.table.AbstractTableModel;
import javax.swing.table.TableColumn;
import javax.swing.table.TableModel;

import beast.base.evolution.alignment.Alignment;
import beast.base.evolution.datatype.DataType;
import beast.base.parser.NexusParser;
import beastfx.app.beauti.ThemeProvider;
import javafx.application.Platform;
import javafx.embed.swing.JFXPanel;
import javafx.embed.swing.SwingNode;
import javafx.scene.control.ButtonType;
import javafx.scene.control.Dialog;
import javafx.scene.control.DialogPane;


// RRB: note that this is still using swing components
// but since this does not seem to require testing, it should be fine.

public class AlignmentViewer extends JPanel {
    private static final long serialVersionUID = 1L;

    Object[][] tableData;
    Object[] columnData;
    boolean useColor = false;
    // flag to indicate that the most frequently occurring character is shown as a dot
    boolean useDots = true;
    Alignment m_alignment;
    Map<Character, Color> m_customColorMap = new HashMap<>();

    /**
     * define which character maps to which color *
     */
    public void setCustomColorMap(Map<Character, Color> colorMap) {
        for (char c : m_customColorMap.keySet()) {
            m_customColorMap.put(c, colorMap.get(c));
        }
    }

    /**
     * constructor processes alignment and sets up table with first column fixed *
     */
    public AlignmentViewer(Alignment data) {
        m_alignment = data;
        int siteCount = data.getSiteCount();
        int taxonCount = data.getTaxonCount();
        tableData = new Object[taxonCount][siteCount + 1];
        char[] headerChar = updateTableData();

        // set up row labels
        for (int i = 0; i < taxonCount; i++) {
            tableData[i][0] = data.getTaxaNames().get(i);
        }

        // set up column labels
        columnData = new Object[siteCount + 1];
        for (int i = 1; i < siteCount + 1; i++) {
            columnData[i] = "<html>.<br>" + headerChar[i - 1] + "</html>";
        }
        columnData[0] = "<html><br>taxon name</html>";
        columnData[1] = "<html>1<br>" + headerChar[0] + "</html>";
        for (int i = 10; i < siteCount; i += 10) {
            String s = i + "";
            for (int j = 0; j < s.length(); j++) {
            	if (i+j < columnData.length) {
            		columnData[i + j] = "<html>" + s.charAt(j) + "<br>" + headerChar[i - 1] + "</html>";
            	}
            }
            columnData[i - 5] = "<html>+<br>" + headerChar[i - 1] + "</html>";
        }

        // create table in scrollpane with first column fixed
        final TableModel fixedColumnModel = new AbstractTableModel() {
            private static final long serialVersionUID = 1L;

            @Override
			public int getColumnCount() {
                return 1;
            }

            @Override
			public String getColumnName(int column) {
                return columnData[column] + "";
            }

            @Override
			public int getRowCount() {
                return tableData.length;
            }

            @Override
			public Object getValueAt(int row, int column) {
                return tableData[row][column];
            }
        };

        final TableModel mainModel = new AbstractTableModel() {
            private static final long serialVersionUID = 1L;

            @Override
			public int getColumnCount() {
                return columnData.length - 1;
            }

            @Override
			public String getColumnName(int column) {
                return columnData[column + 1] + "";
            }

            @Override
			public int getRowCount() {
                return tableData.length;
            }

            @Override
			public Object getValueAt(int row, int column) {
                return tableData[row][column + 1];
            }
        };

        JTable fixedTable = new JTable(fixedColumnModel);
        fixedTable.setAutoResizeMode(JTable.AUTO_RESIZE_OFF);
        Font font = fixedTable.getFont();
        font = new Font(font.getFontName(), font.getStyle(), font.getSize() * 2/3);
        fixedTable.setFont(font);
        TableColumn col = fixedTable.getColumnModel().getColumn(0);
        col.setPreferredWidth(200);
        fixedTable.getTableHeader().setFont(font);

        JTable mainTable = new JTable(mainModel);
        mainTable.setAutoResizeMode(JTable.AUTO_RESIZE_OFF);
        mainTable.setFont(font);
        mainTable.getTableHeader().setFont(font);
        for (int i = 0; i < siteCount; i++) {
            col = mainTable.getColumnModel().getColumn(i);
            col.setPreferredWidth(6);
        }

        ListSelectionModel model = fixedTable.getSelectionModel();
        mainTable.setSelectionModel(model);
        mainTable.setShowGrid(false);

        JScrollPane scrollPane = new JScrollPane(mainTable);
        Dimension fixedSize = fixedTable.getPreferredSize();
        JViewport viewport = new JViewport();
        viewport.setView(fixedTable);
        viewport.setPreferredSize(fixedSize);
        viewport.setMaximumSize(fixedSize);
        scrollPane.setCorner(ScrollPaneConstants.UPPER_LEFT_CORNER, fixedTable.getTableHeader());
        scrollPane.setRowHeaderView(viewport);


        setLayout(new BorderLayout());
        add(scrollPane, BorderLayout.CENTER);
        
        
	      Box buttonBox = Box.createHorizontalBox();
	      JCheckBox useDotsCheckBox = new JCheckBox("Use dots", true);
	      useDotsCheckBox.addActionListener(e -> {
	              JCheckBox _useDots = (JCheckBox) e.getSource();
	              useDots = _useDots.isSelected();
	              updateTableData();
	              repaint();
	          });
	      buttonBox.add(useDotsCheckBox);
	
	      JCheckBox useColorCheckBox = new JCheckBox("Use Color");
	      useColorCheckBox.setName("UseColor");
	      useColorCheckBox.addActionListener(e -> {
	              JCheckBox hasColor = (JCheckBox) e.getSource();
	              useColor = hasColor.isSelected();
	              updateTableData();
	              repaint();
	          });
	      buttonBox.add(useColorCheckBox);
	      add(buttonBox, BorderLayout.SOUTH);
    }

    private char[] updateTableData() {
        int siteCount = m_alignment.getSiteCount();
        int taxonCount = m_alignment.getTaxonCount();

        // set up table content
        DataType dataType = m_alignment.getDataType();
        char[] headerChar = new char[siteCount];
        Object[][] colorMap = setupColorMap();
        Object[][] integerColorMap = setupIntegerColorMap();
        try {
            for (int i = 0; i < siteCount; i++) {
                int patternIndex_ = m_alignment.getPatternIndex(i);
                int[] pattern = m_alignment.getPattern(patternIndex_);
                String patternString = dataType.encodingToString(pattern);

                if (patternString.contains(",")) {
                    // We have a string of comma separated values.
                    String[] patternAtTaxon = patternString.split(",");
                    headerChar[i] = mostFrequentCharInPattern(patternAtTaxon);
                    for (int j = 0; j < taxonCount; j++) {
                        try {
                            int code = Integer.valueOf(patternAtTaxon[j]);
                            if (code + '0' == headerChar[i]) {
                                // Prevent Exception when accessing integerColorMap
                                code = Math.min(Math.max(code, 0), integerColorMap[0].length - 1);
                                tableData[j][i + 1] = integerColorMap[0][code];
                            } else {
                                code = Math.min(Math.max(code, 0), integerColorMap[0].length - 1);
                                tableData[j][i + 1] = integerColorMap[1][code];
                            }
                        } catch (NumberFormatException e) {
                            // State cannot be interpreted as a number.
                            // Assume it is a special (ambiguous) character.
                            tableData[j][i + 1] = useColor ? ' ' : patternAtTaxon[j];
                        }
                    }
                } else {
                    headerChar[i] = mostFrequentCharInPattern(patternString);
                    for (int j = 0; j < taxonCount; j++) {
                        char c = patternString.charAt(j);
                        if (c == headerChar[i]) {
                            tableData[j][i + 1] = colorMap[0][c];
                        } else {
                            tableData[j][i + 1] = colorMap[1][c];
                        }
                    }
                }
            }
        } catch (Exception e) {
            // ignore
        }
        return headerChar;
    }

    /**
     * determine content of table cells.
     * Without color, only Characters are displayed, which can be a bit faster than using color
     * With color, the color is encoded in HTML
     *
     * @return an array of 2x256 where the first entry is for the most frequently occurring character,
     *         and the second for the others
     *         *
     */
    private Object[][] setupColorMap() {
        if (useColor) {
            String[][] colorMap = new String[2][256];
            for (int k = 'A'; k < 'Z'; k++) {
                int i = k - 'A';
                int red = ((i & 0x80) >> 7) + ((i & 0x10) >> 4) + ((i & 0x2) << 1);
                int green = ((i & 0x40) >> 6) + ((i & 0x08) >> 2) + ((i & 0x4));
                int blue = ((i & 0x20) >> 5) + ((i & 0x04) >> 1) + ((i & 0x1) << 2);
                int color = (red << 21 + (green << 18)) + (green << 13) + (blue << 10) + (blue << 5) + (red << 2);
                colorMap[0][k] = "<html><font color='#" + Integer.toString(color, 16) + "'><b>.</b></html>";
                colorMap[1][k] = "<html><font color='#" + Integer.toString(color, 16) + "'><b>" + ((char) k) + "</font></html>";
            }
            for (char c : m_customColorMap.keySet()) {
                Color color = m_customColorMap.get(c);
                colorMap[0][c] = "<html><font color='#" + Integer.toString(color.getRGB(), 16) + "'><b>.</b></html>";
                colorMap[1][c] = "<html><font color='#" + Integer.toString(color.getRGB(), 16) + "'><b>" + c + "</font></html>";
            }
            if (!this.useDots) {
                colorMap[0] = colorMap[1];
            }
            return colorMap;
        } else {
            Character[][] colorMap = new Character[2][256];
            for (int i = 0; i < 256; i++) {
                colorMap[0][i] = '.';
                colorMap[1][i] = (char) i;
            }
            if (!this.useDots) {
                colorMap[0] = colorMap[1];
            }
            return colorMap;
        }
    }

    private Object[][] setupIntegerColorMap() {
        if (useColor) {
            String[][] colorMap = new String[2][256];
            for (int i = 0; i < 256; i++) {
                int red = ((i & 0x80) >> 7) + ((i & 0x10) >> 4) + ((i & 0x2) << 1);
                int green = ((i & 0x40) >> 6) + ((i & 0x08) >> 2) + ((i & 0x4));
                int blue = ((i & 0x20) >> 5) + ((i & 0x04) >> 1) + ((i & 0x1) << 2);
                int color = (red << 21 + (green << 18)) + (green << 13) + (blue << 10) + (blue << 5) + (red << 2);
                colorMap[0][i] = "<html><font color='#" + Integer.toString(color, 16) + "'><b>.</b></html>";
                colorMap[1][i] = "<html><font color='#" + Integer.toString(color, 16) + "'><b>" + i + "</font></html>";
            }
            for (char c : m_customColorMap.keySet()) {
                Color color = m_customColorMap.get(c);
                colorMap[0][c] = "<html><font color='#" + Integer.toString(color.getRGB(), 16) + "'><b>.</b></html>";
                colorMap[1][c] = "<html><font color='#" + Integer.toString(color.getRGB(), 16) + "'><b>" + (int) c + "</font></html>";
            }
            if (!this.useDots) {
                colorMap[0] = colorMap[1];
            }
            return colorMap;
        } else {
            String[][] colorMap = new String[2][256];
            for (int i = 0; i < 256; i++) {
                colorMap[0][i] = ".";
                colorMap[1][i] = Integer.toString(i);
            }
            if (!this.useDots) {
                colorMap[0] = colorMap[1];
            }
            return colorMap;
        }
    }

    private char mostFrequentCharInPattern(String pattern) {
        char[] counts = new char[256];
        for (int i = 0; i < pattern.length(); i++) {
            counts[pattern.charAt(i)]++;
        }
        int maxIndex = 0, max = 0;
        for (int i = 0; i < counts.length; i++) {
            if (counts[i] > max) {
                maxIndex = i;
                max = counts[i];
            }
        }
        return (char) maxIndex;
    }

    private char mostFrequentCharInPattern(String[] pattern) {
        Map<Character, Integer> stateToCount = new HashMap<>();
        int maxCount = 0;
        char mostFrequentChar = 0;
        for (String str : pattern) {
            try {
                char c = (char) (Integer.parseInt(str) + '0');
                if (stateToCount.containsKey(c)) {
                    stateToCount.put(c, stateToCount.get(c) + 1);
                } else {
                    stateToCount.put(c, 1);
                }
                if (stateToCount.get(c) > maxCount) {
                    maxCount = stateToCount.get(c);
                    mostFrequentChar = c;
                }
            } catch (NumberFormatException e) {
                // ignore the ambiguous characters
            }
        }
        return mostFrequentChar;
    }

    public static void showInDialog(Alignment data) {
    	Dialog dialog = new Dialog<>();
    	DialogPane pane = dialog.getDialogPane();
    	pane.getButtonTypes().add(ButtonType.CLOSE);
    	pane.setId("AlignmentViewer");
    	SwingNode panel = new SwingNode();
    	panel.setContent(new AlignmentViewer(data));
    	pane.setContent(panel);
    	pane.setPrefSize(1024, 600);
    	dialog.setResizable(true);
    	ThemeProvider.loadStyleSheet(pane.getScene());
    	dialog.showAndWait();
    }

    public static void main(String[] args) {
        try {
            NexusParser parser = new NexusParser();
            parser.parseFile(new File(args[0]));
            Alignment data = parser.m_alignment;
            AlignmentViewer panel = new AlignmentViewer(data);
            new JFXPanel();
            Platform.runLater(new Runnable() {
				@Override
				public void run() {
		            AlignmentViewer.showInDialog(data);
				}
			});
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
