package fr.perceptron;

import javax.swing.JFrame;
import javax.swing.JPanel;
import java.awt.Color;
import java.awt.Dimension;
import java.awt.Graphics;
import java.awt.event.KeyEvent;
import java.awt.event.KeyListener;
import java.awt.image.BufferedImage;

public class Window {
    
    public Canvas canvas;

    public Window() {}

    public void init(NeuralNetwork nn, DataSet data) {

        // Set the canvas to be the main frame's content pane
        this.canvas = new Canvas(800, 800);
        this.canvas.setBackground(Color.white);
        this.canvas.setPreferredSize(new Dimension(800, 800));
        this.canvas.setVisible(true);
        JFrame mainFrame = new JFrame("Perceptron");
        mainFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        mainFrame.setSize(800, 800);
        mainFrame.add(this.canvas);
        mainFrame.addKeyListener(canvas);
        mainFrame.setResizable(false);
        mainFrame.setLocationRelativeTo(null); // Center the window
        mainFrame.pack();
        mainFrame.setVisible(true);
        mainFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        this.canvas.data = data;
        this.canvas.nn = nn;
        this.canvas.updatePixels();
    }

    public void update() {
        // Update the canvas with the neural network's output
        canvas.updatePixels();
    }
}

class Canvas extends JPanel implements KeyListener {

    private byte[][] pixels;
    private int width;
    private int height;
    public DataSet data;
    private float zoomLevel;
    public NeuralNetwork nn;

    private static final Color[] COLORS = {
        new Color(0xc06699ff, true), // light blue
        new Color(0xc0ff6666, true), // light red
    };

    public Canvas(int width, int height) {
        setBackground(Color.white);
        this.pixels = new byte[height][width];
        this.width = width;
        this.height = height;
        // this.zoomLevel = 2.0f; // linear training set
        // this.zoomLevel = 0.75f; // moon training set
        this.zoomLevel = 0.125f;    // two circles training set

        // Enable double buffering
        setDoubleBuffered(true);
    }

    public Dimension getPreferredSize() {
        return new Dimension(800, 800);
    }

    public void paintComponent(Graphics g) {
        super.paintComponent(g);
        
        // Clear the canvas
        g.setColor(Color.WHITE);
        g.fillRect(0, 0, getWidth(), getHeight());

        // Draw the axes
        this.drawAxis(g);
        // Draw the points
        this.drawDataPoints(g);

        // draw the pixels from the pixels array to the canvas
        // for (int y = 0; y < height; y++) {
        //     for (int x = 0; x < width; x++) {
        //         // Set the color based on the pixel value
        //         g.setColor(COLORS[pixels[y][x]]);
        //         // Draw one pixel
        //         g.fillRect(x, y, 1, 1);
        //     }
        // }

        // Draw the pixels from the pixels array to the canvas
        // use a BufferedImage for better performance
        BufferedImage img = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                // Set the color based on the pixel value
                img.setRGB(x, y, COLORS[pixels[y][x]].getRGB());
            }
        }
        g.drawImage(img, 0, 0, null);
    }

    public void updatePixels() {
        System.out.println("Updating pixels...");

        int midX = width / 2;
        int midY = height / 2;

        // Initialize a new pixels array
        byte[][] newPixels = new byte[height][width];

        // Use parallel processing to compute pixel values
        java.util.stream.IntStream.range(0, height).parallel().forEach(y -> {
            for (int x = 0; x < width; x++) {

                // Normalize the pixel coordinates to the range [-2, 2]
                double[] inputs = {
                    // Normalize the pixel coordinates to the range [-2, 2]
                    (double) (x - midX) / (width / 4 * this.zoomLevel),
                    (double) -(y - midY) / (height / 4 * this.zoomLevel),
                };

                // Calculate the output of the neural network
                double[] output = nn.calculateOutputs(inputs);
                int classification = nn.classify(output[0]);

                // Store the computed value in the temporary array
                newPixels[y][x] = (byte) classification;
            }
        });

        // Replace the old pixels array with the new one
        this.pixels = newPixels;
        System.out.println("Done updating pixels.");

        // Repaint the canvas
        repaint();
    }

    public void drawAxis(Graphics g) {
        // Draw the axes
        g.setColor(Color.BLACK);
        g.drawLine(0, height / 2, width, height / 2); // X-axis
        g.drawLine(width / 2, 0, width / 2, height); // Y-axis

        // Add labels to the axes
        g.drawString("X", width - 20, height / 2 + 15);
        g.drawString("Y", width / 2 - 15, 15);
        // g.drawString("0", width / 2 + 5, height / 2 + 15);
        // g.drawString("1", width - 20, height / 2 - 5);
        // g.drawString("-1", width / 2 + 5, height - 5);
        // g.drawString("-1", 5, height / 2 + 15);
        // g.drawString("1", width / 2 + 5, 15);
    }

    public void drawDataPoints(Graphics g) {
        // Draw the data points
        g.setColor(Color.BLACK);
        for (int i = 0; i < this.data.points.length; i++) {
            int x = (int) (this.zoomLevel * this.data.points[i].inputs[0] * (width / 4) + (width / 2));
            int y = (int) (this.zoomLevel * -this.data.points[i].inputs[1] * (height / 4) + (height / 2));
            if (this.data.points[i].expectedOutputs[0] == 1.0) {
                // Draw a circle
                g.drawOval(x, y, 8, 8);
            } else {
                // Draw a cross
                g.drawLine(x - 4, y, x + 4, y);
                g.drawLine(x, y - 4, x, y + 4);
            }
        }
    }

    @Override
    public void keyPressed(KeyEvent e) {
        
        // Check if the key pressed is UP or DOWN
        if (e.getKeyCode() == KeyEvent.VK_UP) {
            this.zoomLevel += 0.125f;
            System.out.println("Zoom level: " + this.zoomLevel);
        } else if (e.getKeyCode() == KeyEvent.VK_DOWN) {
            this.zoomLevel -= 0.125f;
            System.out.println("Zoom level: " + this.zoomLevel);
        }
        if (this.zoomLevel < 0.125f) {
            this.zoomLevel = 0.125f;
        }
    }

    @Override
    public void keyReleased(KeyEvent e) {}

    @Override
    public void keyTyped(KeyEvent e) {
        
        if (e.getKeyChar() == 'u') {
            // Update the pixels
            updatePixels();
            repaint();
        }
    }
}