import java.awt.BorderLayout;
import java.awt.Dimension;
import java.awt.Graphics;
import java.awt.GridBagConstraints;
import java.awt.GridBagLayout;
import java.awt.Image;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.MouseEvent;
import java.awt.event.MouseListener;
import java.awt.event.MouseMotionListener;
import java.awt.image.MemoryImageSource;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import javax.swing.Box;
import javax.swing.BoxLayout;
import javax.swing.ButtonGroup;
import javax.swing.JButton;
import javax.swing.JComponent;
import javax.swing.JFrame;
import javax.swing.JLabel;
import javax.swing.JOptionPane;
import javax.swing.JPanel;
import javax.swing.JRadioButton;
import javax.swing.JTextField;

public class HeatGui {
    private static final int DEFAULT_WIDTH = 512;
    private static final int DEFAULT_HEIGHT = 512;
    
    public static void main(String[] args) {
        HeatModelManager modelManager = new HeatModelManager();
        modelManager.setModel(new HeatModelSimple(DEFAULT_WIDTH, DEFAULT_HEIGHT, 50.0, 0.2));
        
        FrameContents contents = new FrameContents();

        JFrame frame = new JFrame("2D Heat");
        frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        frame.getContentPane().add(contents.getPanel());
        frame.pack();
        frame.setVisible(true);
        
        Executor executor = new Executor(modelManager, contents);
        executor.start();

        (new HeatEventHandler(modelManager, contents, executor)).register();
    }
    
    private static class FrameContents {
        public final JLabel labelStats;
        public final HeatCanvas canvas;
        public final JRadioButton buttonHot;
        public final JRadioButton buttonCold;
        public final JButton buttonReset;
        
        public FrameContents() {
            canvas = new HeatCanvas(DEFAULT_WIDTH, DEFAULT_HEIGHT);
            labelStats = new JLabel();
            buttonReset = new JButton("Reset...");
            buttonHot = new JRadioButton("Hot", true);
            buttonCold = new JRadioButton("Cold", false);
            ButtonGroup heatGroup = new ButtonGroup();
            heatGroup.add(buttonHot);
            heatGroup.add(buttonCold);
        }
        
        public JPanel getPanel() {
            Box heatPanel = new Box(BoxLayout.PAGE_AXIS);
            heatPanel.setAlignmentX(0.0f);
            heatPanel.add(buttonHot);
            heatPanel.add(buttonCold);
            
            JPanel buttonPanel = new JPanel();
            buttonPanel.add(heatPanel);
            buttonPanel.add(buttonReset);

            JPanel result = new JPanel(new BorderLayout());
            result.add(labelStats, BorderLayout.PAGE_START);
            result.add(canvas, BorderLayout.CENTER);
            result.add(buttonPanel, BorderLayout.PAGE_END);
            
            return result;
        }
    }
    
    private static class Executor {
        private HeatModelManager modelManager;
        private FrameContents frameContents;
        private ScheduledThreadPoolExecutor timer = new ScheduledThreadPoolExecutor(1);
        private Object lock = new Object();
        private ScheduledFuture<?> modelUpdater = null;
        private ScheduledFuture<?> canvasUpdater = null;
        
        public Executor(HeatModelManager modelManager, FrameContents frameContents) {
            this.modelManager = modelManager;
            this.frameContents = frameContents;
        }
        
        public void start() {
            synchronized (lock) {
                if (modelUpdater == null) {
                    modelUpdater = timer.scheduleWithFixedDelay(new ModelUpdater(modelManager),
                            0, 1, TimeUnit.MICROSECONDS);
                }
                if (canvasUpdater == null) {
                    canvasUpdater = timer.scheduleAtFixedRate(new CanvasUpdater(modelManager, frameContents),
                            0, 50, TimeUnit.MILLISECONDS);
                }
            }
        }
        
        public void stop() {
            synchronized (lock) {
                if (modelUpdater != null) {
                    modelUpdater.cancel(false);
                    modelUpdater = null;
                }
                if (canvasUpdater != null) {
                    canvasUpdater.cancel(false);
                    canvasUpdater = null;
                }
            }
        }       
    }
    
    private static class ModelUpdater implements Runnable {
        private HeatModelManager modelManager;
        
        public ModelUpdater(HeatModelManager modelManager) {
            this.modelManager = modelManager;
        }
        
        public void run() {
            modelManager.step();
        }
    }
    
    private static class CanvasUpdater implements Runnable {
        private HeatModelManager modelManager;
        private FrameContents frameContents;
        
        public CanvasUpdater(HeatModelManager modelManager, FrameContents frameContents) {
            this.modelManager = modelManager;
            this.frameContents = frameContents;
        }
        
        public void run() {
            frameContents.canvas.update(modelManager);
            frameContents.labelStats.setText(modelManager.getStatisticsString());
        }
    }
    
    private static class HeatCanvas extends JComponent {
        private static final long serialVersionUID = 1L;
        
        private int width;
        private int height;
        private int[] pix;
        private MemoryImageSource imgSource;
        private Image img = null;
        private double[][] heat;

        public HeatCanvas(int width, int height) {
            this.setPreferredSize(new Dimension(width, height));
            this.width = width;
            this.height = height;
            heat = new double[width][height];
            pix = new int[width * height];
            imgSource = new MemoryImageSource(width, height,
                pix, 0, width);
            imgSource.setAnimated(true);
        }

        public void paintComponent(Graphics g) {
            if (img == null) {
                img = createImage(imgSource);
            }
            g.drawImage(img, 0, 0, this);
        }

        public void update(HeatModelManager modelManager) {
            modelManager.getValues(heat);
            int k = 0;
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    pix[k] = heatToPix(heat[y][x]);
                    k++;
                }
            }
            imgSource.newPixels(0, 0, width, height);
        }

        private static int heatToPix(double heat) {
            int r;
            int g;
            int b;
            if (heat < 50.0) {
                int x = (int) (0.5 + (50.0 - heat) / 50.0 * 127);
                r = 127 - x;
                g = 127 - x;
                b = 127 + x;
            } else {
                int x = (int) (0.5 + (heat - 50.0) / 50.0 * 127);
                r = 127 + x;
                g = 127 - x;
                b = 127 - x;
            }
            return (255 << 24) | (r << 16) | (g << 8) | b;
        }
    }

    private static class HeatEventHandler implements MouseListener, MouseMotionListener, ActionListener {
        private HeatModelManager modelManager;
        private FrameContents frameContents;
        private Executor executor;
        private double pointHeat = 100.0;
        
        public HeatEventHandler(HeatModelManager modelManager, FrameContents frameContents,
                Executor executor) {
            this.modelManager = modelManager;
            this.frameContents = frameContents;
            this.executor = executor;
        }
        
        public void register() {
            frameContents.canvas.addMouseListener(this);
            frameContents.canvas.addMouseMotionListener(this);
            frameContents.buttonReset.addActionListener(this);
            frameContents.buttonHot.addActionListener(this);
            frameContents.buttonCold.addActionListener(this);
        }

        public void mouseDragged(MouseEvent e) {
            mousePressed(e);
        }

        public void mouseMoved(MouseEvent e) { }

        public void mouseClicked(MouseEvent e) { }

        public void mouseEntered(MouseEvent e) {
            if (e.getButton() != MouseEvent.NOBUTTON) {
                mousePressed(e);
            }
        }

        public void mouseExited(MouseEvent e) {
            mouseReleased(e);
        }

        public void mousePressed(MouseEvent e) {
            modelManager.setPinned(e.getX(), e.getY(), pointHeat);
        }

        public void mouseReleased(MouseEvent e) {
            modelManager.clearPinned();
        }
        
        public void actionPerformed(ActionEvent e) {
            if (e.getSource() == frameContents.buttonReset) {
                try {
                    executor.stop();
                    doReset(frameContents.buttonReset, modelManager);
                } finally {
                    executor.start();
                }
            } else {
                pointHeat = frameContents.buttonCold.isSelected() ? 0.0 : 100.0;
            }
        }
    }

    private static void doReset(JComponent parent, HeatModelManager modelManager) {
        String sTemp = "50.0";
        String sConduct = "0.2";
        String sSubWidth = "512";
        String sSubHeight = "512";
        String errMessage = "";
        while (errMessage != null) {
            GridBagLayout gb = new GridBagLayout();
            JPanel panel = new JPanel(gb);
            JTextField fieldTemp = addRow(panel, gb, 0, "Initial Temperature", sTemp);
            JTextField fieldConduct = addRow(panel, gb, 1, "Conductivity", sConduct);
            JTextField fieldSubWidth = addRow(panel, gb, 2, "Subgrid Width", sSubWidth);
            JTextField fieldSubHeight = addRow(panel, gb, 3, "Subgrid Height", sSubHeight);
            int select = JOptionPane.showOptionDialog(parent, panel, "Reset Grid",
                    JOptionPane.OK_CANCEL_OPTION, JOptionPane.PLAIN_MESSAGE,
                    null, null, null);
            if (select != JOptionPane.OK_OPTION) return;
            sTemp = fieldTemp.getText().trim();
            sConduct = fieldConduct.getText().trim();
            sSubWidth = fieldSubWidth.getText().trim();
            sSubHeight = fieldSubHeight.getText().trim();
            errMessage = checkFields(sTemp, sConduct, sSubWidth, sSubHeight);
            if (errMessage != null) {
                select = JOptionPane.showConfirmDialog(parent, errMessage + ".", "Parameter Invalid",
                        JOptionPane.OK_CANCEL_OPTION, JOptionPane.ERROR_MESSAGE);
                if (select != JOptionPane.OK_OPTION) return;
            }
        }
        
        double temp = Double.parseDouble(sTemp);
        double conduct = Double.parseDouble(sConduct);
        int subWidth = Integer.parseInt(sSubWidth);
        int subHeight = Integer.parseInt(sSubHeight);
        HeatModel newModel;
        if (subWidth == DEFAULT_WIDTH && subHeight == DEFAULT_HEIGHT) {
            newModel = new HeatModelSimple(DEFAULT_WIDTH, DEFAULT_HEIGHT, temp, conduct);
        } else {
            newModel = new HeatModelThreaded(DEFAULT_WIDTH, DEFAULT_HEIGHT, temp, conduct,
                    subWidth, subHeight);
        }
        modelManager.setModel(newModel);
    }
    
    private static JTextField addRow(JPanel panel, GridBagLayout gb, int row, String labelText, String init) {
        GridBagConstraints gc = new GridBagConstraints();
        gc.anchor = GridBagConstraints.WEST;
        gc.gridy = row;
        gc.gridx = 0;
        JLabel label = new JLabel(labelText + ": ");
        gb.setConstraints(label, gc);
        panel.add(label);
        gc.gridx = 1;
        gc.fill = GridBagConstraints.HORIZONTAL;
        JTextField field = new JTextField(init);
        field.selectAll();
        gb.setConstraints(field, gc);
        panel.add(field);
        return field;
    }
    
    private static String checkFields(String sTemp, String sConduct, String sSubWidth, String sSubHeight) {
        if (sTemp.equals("")) {
            return "Temperature must be provided";
        } else if (sConduct.equals("")) {
            return "Conductivity must be provided";
        } else if (sSubWidth.equals("")) {
            return "Subgrid width must be provided";
        } else if (sSubHeight.equals("")) {
            return "Subgrid height must be provided";
        }
        
        try {
            double temp = Double.parseDouble(sTemp);
            if (temp < 0.0 || temp > 100.0) {
                return "Temperature must be between 0 and 100";
            }
        } catch (NumberFormatException e) {
            return "Temperature must be valid number";
        }
         
        try {
            double conduct = Double.parseDouble(sConduct);
            if (conduct <= 0.0) {
                return "Conductivity must be positive";
            }
        } catch (NumberFormatException e) {
            return "Conductivity must be valid number";
        }
         
        try {
            int value = Integer.parseInt(sSubWidth);
            if (value <= 0 || value > 1024) {
                return "Subgrid width must be between 1 and 1024";
            } else if ((value & (value - 1)) != 0) {
                return "Subgrid width must be power of 2";
            }
        } catch (NumberFormatException e) {
            return "Subgrid width must be valid integer";
        }
         
        try {
            int value = Integer.parseInt(sSubHeight);
            if (value <= 0 || value > 1024) {
                return "Subgrid height must be between 1 and 1024";
            } else if ((value & (value - 1)) != 0) {
                return "Subgrid height must be power of 2";
            }
        } catch (NumberFormatException e) {
            return "Subgrid height must be valid integer";
        }
        
        return null;
    }

    private static class HeatModelManager {
        private HeatModel model;
        private Object lock;
        private long elapseTotal = 0L;
        private int stepCount = 0;
        
        public HeatModelManager() {
            model = null;
            lock = new Object();
        }
        
        public void setModel(HeatModel value) {
            synchronized (lock) {
                HeatModel old = model;
                if (old != value && old != null) {
                    old.retire();
                }
                model = value;
                elapseTotal = 0L;
                stepCount = 0;
            }
        }
        
        public String getStatisticsString() {
            synchronized (lock) {
                int steps = stepCount;
                if (steps == 0) {
                    return "No steps completed yet";
                } else {
                    return String.format("%.3fms per step (%d steps)", 1e-6 * elapseTotal / steps, steps);
                }
            }
        }
        
        public void step() {
            HeatModel m = model;
            if (m != null) {
                synchronized (lock) {
                    long now = System.nanoTime();
                    m.step();
                    long elapse = System.nanoTime() - now;
                    elapseTotal += elapse;
                    stepCount++;
                }
            }
        }
        
        public void clearPinned() {
            HeatModel m = model;
            if (m != null) {
                synchronized (lock) {
                    m.clearPinned();
                }
            }
        }
        
        public void setPinned(int x, int y, double value) {
            HeatModel m = model;
            if (m != null) {
                synchronized (lock) {
                    m.setPinned(x, y, value);
                }
            }
        }
        
        public void getValues(double[][] result) {
            HeatModel m = model;
            if (m != null) {
                synchronized (lock) {
                    m.getAll(result);
                }
            }
        }
    }
}