package de.mathnat.neuralnetwork.graph;

import javax.swing.*;
import java.awt.*;
import java.awt.geom.Line2D;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class GraphPanel extends JPanel {

    private List<Map.Entry<Double, Double>> testData = new ArrayList<>();
    private List<Map.Entry<Double, Double>> trainData = new ArrayList<>();

    @Override
    protected void paintComponent(Graphics graphics) {
        super.paintComponent(graphics);
        draw(testData, graphics, Color.RED, Color.LIGHT_GRAY);
        draw(trainData, graphics, Color.BLUE, Color.LIGHT_GRAY);
    }

    private void draw(List<Map.Entry<Double, Double>> data, Graphics graphics, Color nodeColor, Color edgeColor) {
        int start = 0;
        if(data.size() > getWidth()) {
            start = data.size() - getWidth();
        }
        for (int i = start; i < data.size(); i++) {
            Map.Entry<Double, Double> currentEntry = data.get(i);
            Map.Entry<Double, Double> nextEntry = i + 1 >= data.size() ? null : data.get(i + 1);
            int x0 = currentEntry.getKey().intValue() - start;
            int y0 = getHeight() - currentEntry.getValue().intValue();
            if (nextEntry != null) {
                graphics.setColor(edgeColor);
                int x1 = nextEntry.getKey().intValue() - start;
                int y1 = getHeight() - nextEntry.getValue().intValue();
                drawLine(graphics, x0, y0, x1, y1);
            }
            graphics.setColor(nodeColor);
            drawCircle(graphics, x0, y0, 6);
        }
    }

    private void drawCircle(Graphics graphics, int x, int y, int radius) {
        graphics.fillOval(x - radius / 2, y - radius / 2, radius, radius);
    }

    private void drawLine(Graphics graphics, int x0, int y0, int x1, int y1) {
        Graphics2D graphics2D = (Graphics2D) graphics;
        graphics2D.setStroke(new BasicStroke(2, BasicStroke.CAP_ROUND, BasicStroke.JOIN_ROUND));
        graphics2D.draw(new Line2D.Double(x0, y0, x1, y1));
    }

    public void clearData() {
        testData.clear();
    }

    public void addTestData(Map.Entry<Double, Double> datum) {
        testData.add(datum);
    }

    public void addTrainData(Map.Entry<Double, Double> datum) {
        trainData.add(datum);
    }
}