import java.awt.BorderLayout;
import java.awt.Color;
import java.awt.Dimension;
import java.awt.Font;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.GridBagConstraints;
import java.awt.GridBagLayout;
import java.awt.GridLayout;
import java.awt.LayoutManager;
import java.awt.color.ColorSpace;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.KeyEvent;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.awt.image.BufferedImage;
import java.awt.image.ColorModel;
import java.awt.image.DataBuffer;
import java.awt.image.SampleModel;
import java.awt.image.WritableRaster;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.imageio.ImageIO;
import javax.swing.BorderFactory;
import javax.swing.Box;
import javax.swing.BoxLayout;
import javax.swing.ImageIcon;
import javax.swing.JButton;
import javax.swing.JDialog;
import javax.swing.JFileChooser;
import javax.swing.JFrame;
import javax.swing.JLabel;
import javax.swing.JMenu;
import javax.swing.JMenuBar;
import javax.swing.JMenuItem;
import javax.swing.JOptionPane;
import javax.swing.JPanel;
import javax.swing.SwingConstants;
import javax.swing.SwingUtilities;
import javax.swing.UIManager;
import javax.swing.border.Border;
import javax.swing.filechooser.FileFilter;
import javax.swing.filechooser.FileNameExtensionFilter;

interface ChannelControlCallback {
	public void setVal(ChannelControl source, int val);
}

class ChannelControl extends JPanel {
	JButton buttonUp = new JButton();
	JButton buttonDown = new JButton();
	JLabel readout = new JLabel("0", SwingConstants.CENTER);
	
	int val;
	
	public void setVal(int i) {
		val = i;
		readout.setText("" + val);
	}
	
	public ChannelControl(Color textCol, final int lowLimit, final int upperLimit, final ChannelControlCallback channelControlCallback) {
		readout.setForeground(textCol);
		readout.setFont(readout.getFont().deriveFont(14f).deriveFont(Font.BOLD));
		this.setBackground(Color.black);
		this.setOpaque(true);
		buttonUp.setText("+");
		buttonUp.setBackground(Color.black);
		buttonUp.setOpaque(true);		
		buttonDown.setText("-");
		buttonDown.setBackground(Color.black);
		buttonDown.setOpaque(true);				
		this.setLayout(new GridLayout(3, 1));
		this.add(readout);
		this.add(buttonUp);
		this.add(buttonDown);
		final ChannelControl channelControl = this;
		buttonUp.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e)
            {
                if (val < upperLimit) {
                	val++;
                	setVal(val);
                	channelControlCallback.setVal(channelControl, val);
                }
            }
        });
		
		buttonDown.addActionListener(new ActionListener() {
            public void actionPerformed(ActionEvent e)
            {
                if (val > lowLimit) {
                	val--;
                	setVal(val);
                	channelControlCallback.setVal(channelControl, val);
                }
            }
        });
	}
}

interface SwatchesCallback {
	public void setColour(Swatch panel);
}

class Swatch extends JPanel {
	int index;
	List<Integer> pixelPositions = new ArrayList<Integer>();
	private Color color;
	Border defaultBorder;
	
	public Swatch(int index, Color c) {
		defaultBorder = BorderFactory.createRaisedBevelBorder();
		setBorder(defaultBorder);
		this.index = index;
		setColour(c);
	}
	
	public void addPixelPos(int pixelPos) {
		pixelPositions.add(pixelPos);
	}
	
	public List<Integer> getPixelPositions() {
		return pixelPositions;
	}
	
	public void setColour(Color c) {
		Color c2 = Swatches.megadriveifyColor(c);
		this.color = c2;
		this.setBackground(c2);
	}
	
	public Color getColor()  {
		return color;
	}

	public void setRed(int val8bit) {
		Color c2 = new Color(val8bit, color.getGreen(), color.getBlue());
		setColour(c2);
	}

	public void setGreen(int val8bit) {
		Color c2 = new Color(color.getRed(), val8bit, color.getBlue());
		setColour(c2);
	}

	public void setBlue(int val8bit) {
		Color c2 = new Color(color.getRed(), color.getGreen(), val8bit);
		setColour(c2);
	}

	public void deselect() {
		this.setBorder(defaultBorder);
	}

	public void select() {
		setBorder(BorderFactory.createCompoundBorder(BorderFactory.createLoweredBevelBorder(), BorderFactory.createLineBorder(Color.white)));
	}
}

class Swatches extends JPanel {
	Swatch[] swatches = new Swatch[64];
	SwatchesCallback swatchesCallback;
	HashMap<Integer, Swatch> pixelPosToSwatch;
	
	public Swatches(final SwatchesCallback swatchesCallback) {
		this.setBackground(Color.black);
		this.setOpaque(true);
		this.swatchesCallback = swatchesCallback;
		setLayout(new GridLayout(2, 32));
		for (int i = 0; i < 64; i++) {
			swatches[i] = new Swatch(i, Color.black);
			add(swatches[i]);
			swatches[i].addMouseListener(new MouseAdapter() { 
		          public void mousePressed(MouseEvent me) {
		        	  Swatch source = (Swatch) me.getSource();
		        	  swatchesCallback.setColour(source);		        	  
		            } 
		          }); 
		}		
	}
	
	private void deselectAll() {
		for (int i = 0; i < 64; i++) {
			swatches[i].deselect();
		}
	}
	
	public static Color megadriveifyColor(Color c) {
		int newRed = convert3bitTo8bit(convert8bitTo3bit(c.getRed()));
		int newGreen = convert3bitTo8bit(convert8bitTo3bit(c.getGreen()));
		int newBlue = convert3bitTo8bit(convert8bitTo3bit(c.getBlue()));
		Color newC = new Color(newRed, newGreen, newBlue);
		return newC;
	}
	
	public static int convert8bitTo3bit(int val8bit) {
		double val3bitD = ((double)val8bit) / (255d/7d);//32d;
		return (int) Math.round(val3bitD);
	}
	
	public static int convert3bitTo8bit(int val3bit) {
		double val = ((double)val3bit) * (255d/7d);
		return (int) Math.round(val);
	}
	
	public void initSwatchesFromPixels(int[] pixels) {
		for (int i = 0; i < 64; i++) {
			swatches[i].setColour(Color.black);
		}			
		pixelPosToSwatch = new HashMap<Integer, Swatch>();
		int foundCols = 0;
		Map<Color, Swatch> colorToSwatch = new HashMap<Color, Swatch>();
		int pixelCount = 0;
		for (int sampleIndex = 0; sampleIndex < pixels.length; sampleIndex+= 4) {
			int r = pixels[sampleIndex];
			int g = pixels[sampleIndex + 1];
			int b = pixels[sampleIndex + 2];
			Color c = new Color(r, g, b);
			Swatch swatch = colorToSwatch.get(c);
			if (swatch == null) {
				swatch = swatches[foundCols++];
				swatch.addPixelPos(sampleIndex);
				swatch.setColour(c);
				colorToSwatch.put(c, swatch);
			} else {
				swatch.addPixelPos(sampleIndex);
			}
			pixelPosToSwatch.put(pixelCount++, swatch);
			pixels[sampleIndex] = convert3bitTo8bit(convert8bitTo3bit(pixels[sampleIndex]));       //red
			pixels[sampleIndex + 1] = convert3bitTo8bit(convert8bitTo3bit(pixels[sampleIndex + 1]));     //green
			pixels[sampleIndex + 2] = convert3bitTo8bit(convert8bitTo3bit(pixels[sampleIndex + 2]));     //blue
			//pixels[i + 3] = 255; //alpha
		}
		swatchesCallback.setColour(swatches[0]);
	}
	
    @Override
    public void paintComponent(Graphics g) {
        super.paintComponent(g);
    }

    public void selectFromPixelPos(int pixelPos) {
    	Swatch swatch = pixelPosToSwatch.get(pixelPos);
	  	swatchesCallback.setColour(swatch);	    	
    }
    
	public void select(Swatch panel) {
		deselectAll();
	  	panel.select();
	}  	
}

class ImageJPanel extends JPanel {

	BufferedImage bi = null;
	
    public ImageJPanel() {
    	this.bi = bi; 
    	this.setBackground(new Color(128, 128, 128));    	
    }
    
    public void loadImage(BufferedImage bi) {
    	this.bi = bi;
    }

    @Override
    public void paintComponent(Graphics g) {
        super.paintComponent(g);
        if (bi != null) {        	       
	        Graphics2D g2d = (Graphics2D) g;
	        g2d.drawImage(this.bi, 0, 0, bi.getWidth(null) * MainJFrame.zoom, bi.getHeight(null) * MainJFrame.zoom, null);
        }
    }    
}

public class MainJFrame extends JFrame implements SwatchesCallback, ChannelControlCallback {

	public static int zoom = 2;
	private ImageJPanel origPanel;
	private ImageJPanel previewPanel;
	private JPanel currentColour;
	private Swatches swatches;
	private ChannelControl redChannelControl;
	private ChannelControl greenChannelControl;
	private ChannelControl blueChannelControl;
	private Swatch currentSwatch;
	private BufferedImage bi;
	private int[] pixels;
	private WritableRaster wr;
	private JPanel controls;
	
	public void setZoom(int zoom) {
		MainJFrame.zoom = zoom;
		repaint();
	}
	
	public void testBitMapping() {
		for (int i = 0; i < 256; i++) {
			int mapping = Swatches.convert8bitTo3bit(i);						
			int mappingBack = Swatches.convert3bitTo8bit(mapping);
			System.out.println(i + " -> " + mapping + " -> " + mappingBack);
		}
		System.out.println();
		for (int i = 0; i < 8; i++) {
			int mapping = Swatches.convert3bitTo8bit(i);	
			int mappingBack = Swatches.convert8bitTo3bit(mapping);
			System.out.println(i + " -> " + mapping + " -> " + mappingBack);
		}		
	}
	
    public MainJFrame() {
//    	testBitMapping();
        try { 
        	UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName()); 
        } catch (Exception e) {
        	System.err.println(e); 
        }
        initUI();
    }   
    
    public final void initUI() {    	
        JMenuBar menubar = new JMenuBar();

        JMenu file = new JMenu("File");
        file.setMnemonic(KeyEvent.VK_F);

        JMenuItem load = new JMenuItem("Load");
        load.setMnemonic(KeyEvent.VK_L);
        file.add(load);        

        final MainJFrame panel = this; 
        
        load.addActionListener(new ActionListener() {

            public void actionPerformed(ActionEvent event) {
            	JFileChooser fileopen = new JFileChooser();
            	fileopen.setCurrentDirectory(new File(System.getProperty("user.dir")));
            	//FileFilter filter = new FileNameExtensionFilter("c files", "c");
            	//fileopen.addChoosableFileFilter(filter);
            	int ret = fileopen.showDialog(panel, "Open file");
            	if (ret == JFileChooser.APPROVE_OPTION) {
            	    File file = fileopen.getSelectedFile();            	   
            		BufferedImage origBi = null;
            		BufferedImage previewBi = null;            
        		    try {
						BufferedImage bi = ImageIO.read(file);
	        		    if (!isValid(bi)) {
	        		    	JOptionPane.showMessageDialog(panel,"Error, invalid input image. Must be 32 bit, but 64 cols max.");								;
	        		    	return;
	        		    }
	        		    origBi = ImageIO.read(file);
	        		    previewBi = ImageIO.read(file);
	        		    initBi(origBi, previewBi);         
	        		    controls.setVisible(true);	        		    
					} catch (IOException e) {
						e.printStackTrace();
					}
            	}
            }
        });        
        
        JMenuItem save = new JMenuItem("Save");
        save.setMnemonic(KeyEvent.VK_S);
        file.add(save);        
        
        save.addActionListener(new ActionListener() {

            public void actionPerformed(ActionEvent event) {
            	JFileChooser filesave = new JFileChooser();
            	filesave.setCurrentDirectory(new File(System.getProperty("user.dir")));
            	FileFilter filter = new FileNameExtensionFilter(".png files", "png");
            	filesave.addChoosableFileFilter(filter);
            	int ret = filesave.showDialog(panel, "Save file");
            	if (ret == JFileChooser.APPROVE_OPTION) {
            	    File file = filesave.getSelectedFile();
            	    try {
            	    	String fileString = file.getPath();
            	    	if (!fileString.toLowerCase().endsWith(".png")) {
            	    		fileString = fileString + ".png";
            	    		file = new File(fileString);
            	    	}
						ImageIO.write(bi, "png", file);
					} catch (IOException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}
            	}
            }
        });         
        
        JMenu view = new JMenu("View");
        view.setMnemonic(KeyEvent.VK_V);

        JMenu zoom = new JMenu("Zoom");
        view.add(zoom);
        
        JMenuItem oneZoom = new JMenuItem("1x");
        JMenuItem twoZoom = new JMenuItem("2x");
        JMenuItem threeZoom = new JMenuItem("3x");
        JMenuItem fourZoom = new JMenuItem("4x");
        
        zoom.add(oneZoom);
        zoom.add(twoZoom);
        zoom.add(threeZoom);
        zoom.add(fourZoom);

        oneZoom.addActionListener(new ActionListener() {

            public void actionPerformed(ActionEvent event) {
            	panel.setZoom(1);                
            }
        });
        
        twoZoom.addActionListener(new ActionListener() {

            public void actionPerformed(ActionEvent event) {
            	panel.setZoom(2);
            }
        });
        
        threeZoom.addActionListener(new ActionListener() {

            public void actionPerformed(ActionEvent event) {
            	panel.setZoom(3);
            }
        });        
        
        fourZoom.addActionListener(new ActionListener() {

            public void actionPerformed(ActionEvent event) {
            	panel.setZoom(4);
            }
        });         

        menubar.add(file);
        menubar.add(view);
        setJMenuBar(menubar);

        setTitle("Megadrive Palette Tweaker");
        setSize(1024, 768);
        setLocationRelativeTo(null);
        setDefaultCloseOperation(EXIT_ON_CLOSE);
        
        this.setLayout(new BorderLayout());        
        JPanel contentPanel = new JPanel();
        contentPanel.setLayout(new GridLayout(1,2));        
        contentPanel.setOpaque(true);
        origPanel = new ImageJPanel();
        previewPanel = new ImageJPanel();
        this.add(contentPanel, BorderLayout.CENTER);
        contentPanel.add(origPanel);
        contentPanel.add(previewPanel);
        controls = new JPanel();
        this.add(controls, BorderLayout.SOUTH);
        controls.setVisible(false);
        //setBackground(Color.black);
        controls.setBackground(Color.black);
        JPanel buttons = new JPanel();
        //buttons.setBackground(Color.black);
        buttons.setOpaque(true);
        controls.setLayout(new BorderLayout());
        
        controls.add(buttons, BorderLayout.WEST);
           
        buttons.setPreferredSize(new Dimension(125, 102));
        
        currentColour = new JPanel();
        currentColour.setBackground(Color.BLUE);
        currentColour.setPreferredSize(new Dimension(125,32));
        buttons.setLayout(new GridBagLayout());
        GridBagConstraints gbc = new GridBagConstraints();
        gbc.fill = GridBagConstraints.HORIZONTAL;
        gbc.anchor = GridBagConstraints.NORTHWEST;
        gbc.gridwidth = 1;
        gbc.gridheight = 1;
        gbc.gridx = 0;
        gbc.gridy = 0;     
        buttons.add(currentColour, gbc);

        JPanel channelControls = new JPanel();
        channelControls.setLayout(new GridLayout(1, 3));
        redChannelControl = new ChannelControl(Color.red, 0, 7, this);
        greenChannelControl = new ChannelControl(Color.green, 0, 7, this);
        blueChannelControl = new ChannelControl(Color.blue, 0, 7, this);
        channelControls.add(redChannelControl);
        channelControls.add(greenChannelControl);
        channelControls.add(blueChannelControl);
        channelControls.setBackground(Color.black);
        gbc.fill = GridBagConstraints.HORIZONTAL;
        gbc.anchor = GridBagConstraints.NORTHWEST;
        gbc.gridwidth = 1;
        gbc.gridheight = 2;
        gbc.gridx = 0;
        gbc.gridy = 1;                
        buttons.add(channelControls, gbc);
        
        swatches = new Swatches(this);
        controls.add(swatches, BorderLayout.CENTER);
    }

    protected boolean isValid(BufferedImage bi) {
    	ColorModel cm = bi.getColorModel();
    	ColorSpace cs = cm.getColorSpace();
    	if ((cs.getType() != 5) ||
    	    (cs.getNumComponents() != 3)) {
    		return false;
    	}
    	WritableRaster wr = bi.getRaster();
    	int[] pixels = new int[(bi.getWidth() * 4) * bi.getHeight()];
    	wr.getPixels(0, 0, bi.getWidth(), bi.getHeight(), pixels);
    	Set<Color> foundColsSet = new HashSet<Color>();
    	int foundCols = 0;
		for (int i = 0; i < pixels.length; i+= 4) {
			int r = pixels[i];
			int g = pixels[i + 1];
			int b = pixels[i + 2];
			Color c = new Color(r, g, b);
			if (!foundColsSet.contains(c)) {
				foundColsSet.add(c);
				foundCols++;
				if (foundCols > 64) {
					return false;
				}
			}
		}
		return true;
	}

	protected boolean initBi(BufferedImage origBi, BufferedImage previewBi) {    	
    	this.bi = previewBi;
    	pixels = new int[(origBi.getWidth()*4)*origBi.getHeight()];
    	wr = bi.getRaster();
		int width = bi.getWidth();
		int height = bi.getHeight();			
		wr.getPixels(0, 0, width, height, pixels);
    	swatches.initSwatchesFromPixels(pixels);
    	wr.setPixels(0, 0, width, height, pixels);
	    origPanel.loadImage(origBi);
	    previewPanel.loadImage(previewBi);
	    previewPanel.addMouseListener(new MouseAdapter() { 
	          public void mousePressed(MouseEvent me) {	        	  
	        	  int x = me.getX() / MainJFrame.zoom;
	        	  int y = me.getY() / MainJFrame.zoom;	        	  
	        	  int pixelPos = (y * bi.getWidth()) + x;
	        	  swatches.selectFromPixelPos(pixelPos);
	            } 
	          });
	    repaint();      
	    return true;
	}

	public static void main(String[] args) {

        SwingUtilities.invokeLater(new Runnable() {

            public void run() {
            	MainJFrame sd = new MainJFrame();
                sd.setVisible(true);
            }
        });
    }

	@Override
	public void setColour(Swatch panel) {
	  	swatches.select(panel);		
		Color c = panel.getColor();
		int r = Swatches.convert8bitTo3bit(c.getRed());
		int g = Swatches.convert8bitTo3bit(c.getGreen());
		int b = Swatches.convert8bitTo3bit(c.getBlue());
		redChannelControl.setVal(r);
		greenChannelControl.setVal(g);
		blueChannelControl.setVal(b);		
		currentColour.setBackground(panel.getBackground());
		currentSwatch = panel;
	}

	@Override
	public void setVal(ChannelControl source, int val) {
		int val8bit = Swatches.convert3bitTo8bit(val);
		if (source == redChannelControl) {
			currentSwatch.setRed(val8bit);
		} else if (source == greenChannelControl) {
			currentSwatch.setGreen(val8bit);
		} else if (source == blueChannelControl) {
			currentSwatch.setBlue(val8bit);
		}
		Color currentCol = currentSwatch.getColor();
		currentColour.setBackground(currentCol);
		int red = currentCol.getRed();
		int green = currentCol.getGreen();
		int blue = currentCol.getBlue();
		for (int pixelPos : currentSwatch.getPixelPositions()) {
			pixels[pixelPos] = red;
			pixels[pixelPos + 1] = green;
			pixels[pixelPos + 2] = blue;
		}				
		wr.setPixels(0, 0, bi.getWidth(), bi.getHeight(), pixels);		
		repaint();
	}
}

