package sparseviewer.ui;

import sparseviewer.model.PlotMode;
import sparseviewer.interp.InterpolationResult;
import sparseviewer.interp.Grid2D;


import javax.swing.JPanel;

import java.awt.BasicStroke;
import java.awt.Color;
import java.awt.FontMetrics;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.image.BufferedImage;

import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.awt.event.MouseMotionAdapter;

public class PlotPanel extends JPanel {

    private double[] x;
    private double[] y;
    private double[] z;

    private InterpolationResult interpolationResult;
    private BufferedImage imageBuffer;
    private String xName = "X";
    private String yName = "Y";
    private String zName = "Z";

    private PlotMode mode = PlotMode.PLOT_1D_LINE_POINTS;

    private double xmin;
    private double xmax;
    private double ymin;
    private double ymax;
    private double zmin;
    private double zmax;

    private boolean hasData = false;
    private boolean is2D = false;

    private double crossX = Double.NaN;
    private double crossY = Double.NaN;

    private boolean dragVertical = false;
    private boolean dragHorizontal = false;

    private final int leftMargin = 80;
    private final int rightMargin = 45;
    private final int topMargin = 40;
    private final int bottomMargin = 70;

    public PlotPanel() {
        setBackground(Color.WHITE);

        addMouseListener(new MouseAdapter() {
            @Override
            public void mousePressed(MouseEvent e) {
                if (!hasData) {
                    return;
                }

                int px = dataToScreenX(crossX);
                int py = dataToScreenY(crossY);

                if (Math.abs(e.getX() - px) < 8) {
                    dragVertical = true;
                }

                if (Math.abs(e.getY() - py) < 8) {
                    dragHorizontal = true;
                }

                if (!dragVertical && !dragHorizontal) {
                    crossX = screenToDataX(e.getX());
                    crossY = screenToDataY(e.getY());
                    dragVertical = true;
                    dragHorizontal = true;
                }

                repaint();
            }

            @Override
            public void mouseReleased(MouseEvent e) {
                dragVertical = false;
                dragHorizontal = false;
            }
        });

        addMouseMotionListener(new MouseMotionAdapter() {
            @Override
            public void mouseDragged(MouseEvent e) {
                if (!hasData) {
                    return;
                }

                if (dragVertical) {
                    crossX = screenToDataX(e.getX());
                }

                if (dragHorizontal) {
                    crossY = screenToDataY(e.getY());
                }

                repaint();
            }
        });
    }

    public void clear() {
        hasData = false;
        x = null;
        y = null;
        z = null;
        interpolationResult = null;
        imageBuffer = null;
        repaint();
    }

    public void set1DData(double[] x,
                          double[] y,
                          String xName,
                          String yName,
                          PlotMode mode) {
        this.x = x;
        this.y = y;
        this.z = null;
        this.xName = xName;
        this.yName = yName;
        this.zName = "";
        this.mode = mode;
        this.is2D = false;
        this.hasData = x != null
                && y != null
                && x.length > 0
                && x.length == y.length;

        computeBounds();

        crossX = 0.5 * (xmin + xmax);
        crossY = 0.5 * (ymin + ymax);

        repaint();
    }

    public void set2DData(double[] x,
                          double[] y,
                          double[] z,
                          String xName,
                          String yName,
                          String zName) {
        this.x = x;
        this.y = y;
        this.z = z;
        this.xName = xName;
        this.yName = yName;
        this.zName = zName;
        this.mode = PlotMode.PLOT_2D_SCATTER;

        this.interpolationResult = null;
        this.imageBuffer = null;

        this.is2D = true;
        this.hasData = x != null
                && y != null
                && z != null
                && x.length > 0
                && x.length == y.length
                && x.length == z.length;

        computeBounds();

        crossX = 0.5 * (xmin + xmax);
        crossY = 0.5 * (ymin + ymax);

        repaint();
    }

    private void computeBounds() {
        xmin = Double.POSITIVE_INFINITY;
        xmax = Double.NEGATIVE_INFINITY;
        ymin = Double.POSITIVE_INFINITY;
        ymax = Double.NEGATIVE_INFINITY;
        zmin = Double.POSITIVE_INFINITY;
        zmax = Double.NEGATIVE_INFINITY;

        if (!hasData) {
            return;
        }

        for (int i = 0; i < x.length; i++) {
            if (isFinite(x[i])) {
                if (x[i] < xmin) {
                    xmin = x[i];
                }

                if (x[i] > xmax) {
                    xmax = x[i];
                }
            }

            if (isFinite(y[i])) {
                if (y[i] < ymin) {
                    ymin = y[i];
                }

                if (y[i] > ymax) {
                    ymax = y[i];
                }
            }

            if (is2D && z != null && isFinite(z[i])) {
                if (z[i] < zmin) {
                    zmin = z[i];
                }

                if (z[i] > zmax) {
                    zmax = z[i];
                }
            }
        }

        if (!isFinite(xmin) || !isFinite(xmax) || !isFinite(ymin) || !isFinite(ymax)) {
            hasData = false;
            return;
        }

        if (xmin == xmax) {
            xmin -= 1.0;
            xmax += 1.0;
        }

        if (ymin == ymax) {
            ymin -= 1.0;
            ymax += 1.0;
        }

        if (is2D) {
            if (!isFinite(zmin) || !isFinite(zmax)) {
                zmin = 0.0;
                zmax = 1.0;
            }

            if (zmin == zmax) {
                zmin -= 1.0;
                zmax += 1.0;
            }
        }

        double xpad = 0.03 * (xmax - xmin);
        double ypad = 0.03 * (ymax - ymin);

        xmin -= xpad;
        xmax += xpad;
        ymin -= ypad;
        ymax += ypad;
    }

    @Override
    protected void paintComponent(Graphics g) {
        super.paintComponent(g);

        Graphics2D g2 = (Graphics2D) g;
        g2.setRenderingHint(
                RenderingHints.KEY_ANTIALIASING,
                RenderingHints.VALUE_ANTIALIAS_ON
        );

        drawAxes(g2);

        if (!hasData) {
            g2.setColor(Color.GRAY);
            g2.drawString("Load TXT file(s) to start.", leftMargin + 20, topMargin + 30);
            return;
        }

        if (is2D) {
            if (mode == PlotMode.PLOT_2D_IMAGE) {
                drawInterpolatedImage(g2);
            } else {
                drawScatter2D(g2);
            }
        } else {
            drawPlot1D(g2);
        }

        drawCrosshair(g2);
        drawReadout(g2);
    }

    private void drawAxes(Graphics2D g2) {
        int w = getWidth();
        int h = getHeight();

        int x0 = leftMargin;
        int y0 = h - bottomMargin;
        int x1 = w - rightMargin;
        int y1 = topMargin;

        g2.setColor(Color.BLACK);
        g2.drawLine(x0, y0, x1, y0);
        g2.drawLine(x0, y0, x0, y1);

        g2.drawString(shortName(xName), x0 + 10, h - 25);
        g2.drawString(shortName(yName), 15, y1 + 15);

        if (is2D) {
            g2.drawString("Color: " + shortName(zName), x0 + 10, 22);
        }
    }

    private void drawPlot1D(Graphics2D g2) {
        boolean drawLine = mode == PlotMode.PLOT_1D_LINE
                || mode == PlotMode.PLOT_1D_LINE_POINTS;

        boolean drawPoints = mode == PlotMode.PLOT_1D_POINTS
                || mode == PlotMode.PLOT_1D_LINE_POINTS;

        g2.setColor(new Color(30, 90, 180));
        g2.setStroke(new BasicStroke(1.3f));

        int prevX = 0;
        int prevY = 0;
        boolean havePrev = false;

        for (int i = 0; i < x.length; i++) {
            if (!isFinite(x[i]) || !isFinite(y[i])) {
                continue;
            }

            int sx = dataToScreenX(x[i]);
            int sy = dataToScreenY(y[i]);

            if (drawLine && havePrev) {
                g2.drawLine(prevX, prevY, sx, sy);
            }

            if (drawPoints) {
                g2.fillOval(sx - 3, sy - 3, 6, 6);
            }

            prevX = sx;
            prevY = sy;
            havePrev = true;
        }
    }

    private void drawScatter2D(Graphics2D g2) {
        for (int i = 0; i < x.length; i++) {
            if (!isFinite(x[i]) || !isFinite(y[i]) || !isFinite(z[i])) {
                continue;
            }

            int sx = dataToScreenX(x[i]);
            int sy = dataToScreenY(y[i]);

            g2.setColor(colorForValue(z[i]));
            g2.fillOval(sx - 3, sy - 3, 6, 6);
        }

        drawColorScale(g2);
    }

    private void drawColorScale(Graphics2D g2) {
        int w = getWidth();
        int h = getHeight();

        int x0 = w - 30;
        int y0 = topMargin;
        int barH = h - topMargin - bottomMargin;

        for (int i = 0; i < barH; i++) {
            double t = 1.0 - (double) i / Math.max(1, barH - 1);
            double val = zmin + t * (zmax - zmin);

            g2.setColor(colorForValue(val));
            g2.drawLine(x0, y0 + i, x0 + 12, y0 + i);
        }

        g2.setColor(Color.BLACK);
        g2.drawRect(x0, y0, 12, barH);

        g2.drawString(format(zmax), x0 - 65, y0 + 10);
        g2.drawString(format(zmin), x0 - 65, y0 + barH);
    }

    private void drawCrosshair(Graphics2D g2) {
        if (!isFinite(crossX) || !isFinite(crossY)) {
            return;
        }

        int sx = dataToScreenX(crossX);
        int sy = dataToScreenY(crossY);

        int w = getWidth();
        int h = getHeight();

        g2.setColor(new Color(0, 0, 0, 150));
        g2.setStroke(new BasicStroke(1.0f));

        g2.drawLine(sx, topMargin, sx, h - bottomMargin);
        g2.drawLine(leftMargin, sy, w - rightMargin, sy);
    }

    private void drawReadout(Graphics2D g2) {
        int h = getHeight();

        String text = "X=" + format(crossX) + "   Y=" + format(crossY);

        if (is2D) {
            if (mode == PlotMode.PLOT_2D_IMAGE && interpolationResult != null) {
                float value = interpolationResult.getNearestValue(crossX, crossY);

                text += "   Z(interp)=" + format(value);

                int idx = nearestPoint(crossX, crossY);
                if (idx >= 0 && z != null) {
                    text += "   nearest raw Z=" + format(z[idx])
                            + "   idx=" + idx;
                }
            } else {
                int idx = nearestPoint(crossX, crossY);

                if (idx >= 0) {
                    text += "   nearest Z=" + format(z[idx])
                            + "   idx=" + idx;
                }
            }
        } else {
            int idx = nearestX(crossX);

            if (idx >= 0) {
                text += "   nearest Y=" + format(y[idx])
                        + "   idx=" + idx;
            }
        }

        g2.setColor(Color.BLACK);
        g2.drawString(text, leftMargin + 10, h - 8);
    }
    private void drawInterpolatedImage(Graphics2D g2) {
        if (imageBuffer == null) {
            return;
        }

        int x0 = leftMargin;
        int y0 = topMargin;
        int w = getWidth() - leftMargin - rightMargin;
        int h = getHeight() - topMargin - bottomMargin;

        g2.drawImage(imageBuffer, x0, y0, w, h, null);

        g2.setColor(Color.BLACK);
        g2.drawRect(x0, y0, w, h);

        drawColorScale(g2);
    }
    private int nearestX(double xv) {
        int best = -1;
        double bestD = Double.POSITIVE_INFINITY;

        for (int i = 0; i < x.length; i++) {
            if (!isFinite(x[i])) {
                continue;
            }

            double d = Math.abs(x[i] - xv);

            if (d < bestD) {
                bestD = d;
                best = i;
            }
        }

        return best;
    }

    private int nearestPoint(double xv, double yv) {
        int best = -1;
        double bestD = Double.POSITIVE_INFINITY;

        double xr = xmax - xmin;
        double yr = ymax - ymin;

        for (int i = 0; i < x.length; i++) {
            if (!isFinite(x[i]) || !isFinite(y[i])) {
                continue;
            }

            double dx = (x[i] - xv) / xr;
            double dy = (y[i] - yv) / yr;
            double d = dx * dx + dy * dy;

            if (d < bestD) {
                bestD = d;
                best = i;
            }
        }

        return best;
    }

    private int dataToScreenX(double xv) {
        int w = getWidth();
        double t = (xv - xmin) / (xmax - xmin);

        return leftMargin + (int) Math.round(t * (w - leftMargin - rightMargin));
    }

    private int dataToScreenY(double yv) {
        int h = getHeight();
        double t = (yv - ymin) / (ymax - ymin);

        return h - bottomMargin
                - (int) Math.round(t * (h - topMargin - bottomMargin));
    }
    private void rebuildImageBuffer() {
        if (interpolationResult == null) {
            imageBuffer = null;
            return;
        }

        Grid2D grid = interpolationResult.getGrid();
        int width = grid.getWidth();
        int height = grid.getHeight();

        imageBuffer = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        float[] pixels = interpolationResult.getPixels();

        for (int iy = 0; iy < height; iy++) {
            for (int ix = 0; ix < width; ix++) {
                float v = pixels[iy * width + ix];

                Color c;

                if (Float.isNaN(v) || Float.isInfinite(v)) {
                    c = Color.GRAY;
                } else {
                    c = colorForValue(v);
                }

                imageBuffer.setRGB(ix, height - 1 - iy, c.getRGB());
            }
        }
    }

    public void set2DImageData(double[] x,
                               double[] y,
                               double[] z,
                               String xName,
                               String yName,
                               String zName,
                               InterpolationResult result) {
        this.x = x;
        this.y = y;
        this.z = z;
        this.xName = xName;
        this.yName = yName;
        this.zName = zName;
        this.mode = PlotMode.PLOT_2D_IMAGE;
        this.is2D = true;
        this.hasData = result != null;
        this.interpolationResult = result;

        if (result != null) {
            Grid2D grid = result.getGrid();

            this.xmin = grid.getXMin();
            this.xmax = grid.getXMax();
            this.ymin = grid.getYMin();
            this.ymax = grid.getYMax();
            this.zmin = result.getZMin();
            this.zmax = result.getZMax();

            this.crossX = 0.5 * (xmin + xmax);
            this.crossY = 0.5 * (ymin + ymax);

            rebuildImageBuffer();
        }

        repaint();
    }

    private double screenToDataX(int sx) {
        int w = getWidth();

        double t = (double) (sx - leftMargin)
                / Math.max(1, w - leftMargin - rightMargin);

        return xmin + t * (xmax - xmin);
    }

    private double screenToDataY(int sy) {
        int h = getHeight();

        double t = (double) (h - bottomMargin - sy)
                / Math.max(1, h - topMargin - bottomMargin);

        return ymin + t * (ymax - ymin);
    }

    private Color colorForValue(double v) {
        double t = (v - zmin) / (zmax - zmin);

        if (t < 0) {
            t = 0;
        }

        if (t > 1) {
            t = 1;
        }

        float hue = (float) (0.66 - 0.66 * t);

        return Color.getHSBColor(hue, 1.0f, 1.0f);
    }

    private boolean isFinite(double v) {
        return !Double.isNaN(v) && !Double.isInfinite(v);
    }

    private String format(double v) {
        if (!isFinite(v)) {
            return "NaN";
        }

        return String.format("%.6g", v);
    }

    private String shortName(String name) {
        if (name == null) {
            return "";
        }

        String s = name;

        int typeIdx = s.lastIndexOf(":float");

        if (typeIdx > 0) {
            s = s.substring(0, typeIdx);
        }

        int idx = s.lastIndexOf(':');

        if (idx >= 0 && idx < s.length() - 1) {
            return s.substring(idx + 1);
        }

        return s;
    }
}