package com.naaturel.ANN.infrastructure.visualization;

import com.naaturel.ANN.domain.abstraction.Model;
import com.naaturel.ANN.domain.model.neuron.Neuron;

import javax.swing.*;
import java.awt.*;
import java.util.*;
import java.util.List;

public class ModelVisualizer {

    private final JFrame frame;
    private final Model model;
    private boolean withWeights;

    public ModelVisualizer(Model model){
        this.frame = initFrame();
        this.model = model;
        this.withWeights = false;
    }

    public ModelVisualizer withWeights(boolean value){
        this.withWeights = value;
        return this;
    }

    public void display() {
        JPanel panel = buildPanel();
        frame.add(panel);
        frame.setVisible(true);
    }

    private JFrame initFrame(){
        JFrame frame = new JFrame("Model Visualizer");
        frame.setDefaultCloseOperation(JFrame.DISPOSE_ON_CLOSE);
        frame.setSize(800, 600);
        return frame;
    }

    private Map<Integer, Point> computeNeuronPositions(int width, int height) {

        Map<Integer, List<Neuron>> neuronsByLayer = new LinkedHashMap<>();
        model.forEachNeuron(n -> {
            int layerIndex = model.layerIndexOf(n);
            neuronsByLayer.computeIfAbsent(layerIndex, k -> new ArrayList<>()).add(n);
        });

        Map<Integer, Point> neuronPositions = new HashMap<>();
        int layerCount = neuronsByLayer.size();
        int layerX = width / (layerCount + 1);

        for (Map.Entry<Integer, List<Neuron>> entry : neuronsByLayer.entrySet()) {
            int layerIndex = entry.getKey();
            List<Neuron> neurons = entry.getValue();
            int x = layerX * (layerIndex + 1);
            int neuronCount = neurons.size();

            for (int i = 0; i < neuronCount; i++) {
                int y = height / (neuronCount + 1) * (i + 1);
                neuronPositions.put(neurons.get(i).getId(), new Point(x, y));
            }
        }
        return neuronPositions;
    }

    private void drawConnections(Graphics2D g2, Map<Integer, Point> neuronPositions) {
        model.forEachNeuron(n -> {
            Point from = neuronPositions.get(n.getId());
            int neuronIndex = model.indexInLayerOf(n);
            model.forEachNeuronConnectedTo(n, connected -> {
                Point to = neuronPositions.get(connected.getId());
                g2.setColor(Color.LIGHT_GRAY);
                g2.drawLine(from.x, from.y, to.x, to.y);

                if(!this.withWeights) return;
                int mx = (from.x + to.x) / 2;
                int my = (from.y + to.y) / 2;
                float weight = connected.getWeight(neuronIndex + 1);
                g2.setColor(Color.DARK_GRAY);
                g2.drawString(String.format("%.2f", weight), mx, my);
            });
        });
    }

    private void drawNeurons(Graphics2D g2, Map<Integer, Point> neuronPositions, int neuronRadius) {
        model.forEachNeuron(n -> {
            Point p = neuronPositions.get(n.getId());
            g2.setColor(Color.WHITE);
            g2.fillOval(p.x - neuronRadius, p.y - neuronRadius, neuronRadius * 2, neuronRadius * 2);
            g2.setColor(Color.BLACK);
            g2.drawOval(p.x - neuronRadius, p.y - neuronRadius, neuronRadius * 2, neuronRadius * 2);
            g2.drawString(String.valueOf(n.getId()), p.x - 5, p.y + 5);
        });
    }

    private JPanel buildPanel() {
        int neuronRadius = 20;
        return new JPanel() {
            @Override
            protected void paintComponent(Graphics g) {
                super.paintComponent(g);
                Graphics2D g2 = (Graphics2D) g;
                g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);

                Map<Integer, Point> neuronPositions = computeNeuronPositions(getWidth(), getHeight());
                drawConnections(g2, neuronPositions);
                drawNeurons(g2, neuronPositions, neuronRadius);
            }
        };
    }


}
