import javax.swing.*;
import javax.swing.border.TitledBorder;
import java.awt.*;
import java.io.*;
import java.util.ArrayList;
import java.util.Scanner;

public class AntSimulationGUI extends JFrame {
    private AntSimulationController controller;
    private GridPanel gridPanel;
    private Timer simulationTimer;
    
    private JButton playButton;
    private JButton pauseButton;
    private JButton stepButton;
    private JButton loadMapButton;
    private JSlider speedSlider;
    private JSlider decaySlider;
    private TitledBorder speedBorder;
    private TitledBorder decayBorder;

    private JButton addAntButton;
    private JButton removeAntButton;
    private JLabel antCountLabel;
    
    private JLabel tripCountLabel;
    private int successfulTrips = 0;

    private int gridWidth = 4;
    private int gridHeight = 3;
    private double decayFactor = 0.05; 

    private static final int WINDOW_WIDTH = 1000;
    private static final int WINDOW_HEIGHT = 700;

    public AntSimulationGUI() {
        controller = new AntSimulationController(gridWidth, gridHeight, 1);

        setTitle("CS1 Ant Foraging Simulation");
        setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setSize(WINDOW_WIDTH, WINDOW_HEIGHT);
        setLayout(new BorderLayout());

        gridPanel = new GridPanel();
        add(gridPanel, BorderLayout.CENTER);

        setupControlTower();    
        setupStatusDashboard(); 

        simulationTimer = new Timer(100, e -> tickSimulation());

        setLocationRelativeTo(null);
        setVisible(true);
    }

    private void loadMapFromTextFile(String filename) {
        File file = new File(filename);
        ArrayList<String[]> lines = new ArrayList<>();
        try (Scanner scanner = new Scanner(file)) {
            while (scanner.hasNextLine()) {
                String line = scanner.nextLine().trim();
                if (!line.isEmpty()) {
                    lines.add(line.split("\\s+"));
                }
            }
        } catch (FileNotFoundException e) {
            JOptionPane.showMessageDialog(this, "Map file could not be found.", "Error", JOptionPane.ERROR_MESSAGE);
            return;
        }

        gridHeight = lines.size();
        gridWidth = lines.get(0).length;

        int totalAnts = (controller != null) ? controller.getAnts().size() : 1; 
        controller = new AntSimulationController(gridWidth, gridHeight, totalAnts);
        successfulTrips = 0; 
        tripCountLabel.setText("0");
        gridPanel.repaint();
    }

    private void setupControlTower() {
        JPanel sidePanel = new JPanel();
        sidePanel.setLayout(new BoxLayout(sidePanel, BoxLayout.Y_AXIS));
        sidePanel.setBorder(BorderFactory.createEmptyBorder(15, 15, 15, 15));

        JPanel actionPanel = new JPanel(new GridLayout(4, 1, 0, 8));
        actionPanel.setBorder(BorderFactory.createTitledBorder("Controls"));
        playButton = new JButton("Play");
        pauseButton = new JButton("Pause");
        stepButton = new JButton("Step Forward");
        loadMapButton = new JButton("Load Map");
        pauseButton.setEnabled(false);
        
        actionPanel.add(playButton);
        actionPanel.add(pauseButton);
        actionPanel.add(stepButton);
        actionPanel.add(loadMapButton);
        sidePanel.add(actionPanel);

        sidePanel.add(Box.createVerticalStrut(15));

        JPanel speedPanel = new JPanel(new BorderLayout());
        speedSlider = new JSlider(JSlider.HORIZONTAL, 1, 40, 10);
        speedSlider.setPreferredSize(new Dimension(220, 45));
        speedBorder = BorderFactory.createTitledBorder("Simulation Speed: 10 Hz");
        speedPanel.setBorder(speedBorder);
        speedPanel.add(speedSlider, BorderLayout.CENTER);
        sidePanel.add(speedPanel);

        sidePanel.add(Box.createVerticalStrut(15));

        JPanel decayPanel = new JPanel(new BorderLayout());
        decaySlider = new JSlider(JSlider.HORIZONTAL, 0, 100, (int)(decayFactor * 100));
        decaySlider.setPreferredSize(new Dimension(220, 45)); 
        decayBorder = BorderFactory.createTitledBorder("Decay Rate: 5%");
        decayPanel.setBorder(decayBorder);
        decayPanel.add(decaySlider, BorderLayout.CENTER);
        sidePanel.add(decayPanel);

        playButton.addActionListener(e -> {
            simulationTimer.start();
            playButton.setEnabled(false);
            pauseButton.setEnabled(true);
            stepButton.setEnabled(false);
            loadMapButton.setEnabled(false);
        });

        pauseButton.addActionListener(e -> {
            simulationTimer.stop();
            playButton.setEnabled(true);
            pauseButton.setEnabled(false);
            stepButton.setEnabled(true);
            loadMapButton.setEnabled(true);
        });

        stepButton.addActionListener(e -> tickSimulation());

        loadMapButton.addActionListener(e -> {
            JFileChooser fileChooser = new JFileChooser(".");
            fileChooser.setDialogTitle("Select Map Configuration File");
            if (fileChooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
                loadMapFromTextFile(fileChooser.getSelectedFile().getAbsolutePath());
                antCountLabel.setText(String.valueOf(controller.getAnts().size()));
                gridPanel.repaint();
            }
        });

        speedSlider.addChangeListener(e -> {
            int ticksPerSecond = speedSlider.getValue();
            speedBorder.setTitle("Simulation Speed: " + ticksPerSecond + " Hz");
            speedSlider.repaint();
            sidePanel.revalidate();
            sidePanel.repaint();
            simulationTimer.setDelay(1000 / ticksPerSecond);
        });

        decaySlider.addChangeListener(e -> {
            int currentDecayPercent = decaySlider.getValue();
            decayBorder.setTitle("Decay Rate: " + currentDecayPercent + "%");
            decaySlider.repaint();
            sidePanel.revalidate();
            sidePanel.repaint();
            decayFactor = currentDecayPercent / 100.0;
        });

        add(sidePanel, BorderLayout.EAST);
    }

    private void setupStatusDashboard() {
        JPanel statusPanel = new JPanel(new FlowLayout(FlowLayout.CENTER, 50, 10));
        statusPanel.setBorder(BorderFactory.createEmptyBorder(5, 5, 15, 5));

        JPanel populationPanel = new JPanel(new FlowLayout(FlowLayout.CENTER, 8, 5));
        populationPanel.setBorder(BorderFactory.createTitledBorder("Colony Size"));
        removeAntButton = new JButton("–");
        addAntButton = new JButton("+");
        antCountLabel = new JLabel(String.valueOf(controller.getAnts().size()));
        antCountLabel.setFont(new Font("SansSerif", Font.BOLD, 13));
        antCountLabel.setPreferredSize(new Dimension(50, 20)); 
        antCountLabel.setHorizontalAlignment(SwingConstants.CENTER);
        
        populationPanel.add(removeAntButton);
        populationPanel.add(antCountLabel);
        populationPanel.add(addAntButton);
        statusPanel.add(populationPanel);

        JPanel metricsPanel = new JPanel(new FlowLayout(FlowLayout.CENTER, 15, 5));
        metricsPanel.setBorder(BorderFactory.createTitledBorder("Total Trips Completed"));
        tripCountLabel = new JLabel(String.valueOf(successfulTrips));
        tripCountLabel.setFont(new Font("SansSerif", Font.BOLD, 15));
        tripCountLabel.setForeground(Color.GREEN.darker());
        
        tripCountLabel.setPreferredSize(new Dimension(150, 20)); 
        tripCountLabel.setHorizontalAlignment(SwingConstants.CENTER);
        metricsPanel.add(tripCountLabel);
        statusPanel.add(metricsPanel);

        addAntButton.addActionListener(e -> {
            ArrayList<Ant> ants = controller.getAnts();
            ants.add(new Ant(controller.getNestX(), controller.getNestY(), controller.getGrid()));
            antCountLabel.setText(String.valueOf(ants.size()));
            gridPanel.repaint();
        });

        removeAntButton.addActionListener(e -> {
            ArrayList<Ant> ants = controller.getAnts();
            if (ants.size() > 1) {
                ants.remove(ants.size() - 1);
                antCountLabel.setText(String.valueOf(ants.size()));
                gridPanel.repaint();
            }
        });

        add(statusPanel, BorderLayout.SOUTH);
    }

    private void tickSimulation() {
        ArrayList<Ant> ants = controller.getAnts();
        
        boolean[] antHadFoodSnapshot = new boolean[ants.size()];
        for (int i = 0; i < ants.size(); i++) {
            antHadFoodSnapshot[i] = ants.get(i).hasFood();
        }

        controller.updateSimulation(decayFactor);

        for (int i = 0; i < ants.size(); i++) {
            if (i < antHadFoodSnapshot.length) {
                if (antHadFoodSnapshot[i] && !ants.get(i).hasFood()) {
                    successfulTrips++;
                }
            }
        }
        
        tripCountLabel.setText(String.valueOf(successfulTrips));
        gridPanel.repaint(); 
    }

    private class GridPanel extends JPanel {
        private static final int MAX_TILE_SIZE = 60; 

        public GridPanel() {
            setBackground(new Color(218, 185, 143)); 
            setPreferredSize(new Dimension(650, 550)); 
        }

        @Override
        protected void paintComponent(Graphics g) {
            super.paintComponent(g);
            Graphics2D g2 = (Graphics2D) g;
            g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);

            AntGrid grid = controller.getGrid();
            int cols = grid.getWidth();
            int rows = grid.getHeight();
            
            int padding = 40;
            int tileSize = Math.min((getWidth() - padding) / cols, (getHeight() - padding) / rows);
            if (tileSize > MAX_TILE_SIZE) tileSize = MAX_TILE_SIZE;
            if (tileSize < 12) tileSize = 12; 

            int totalGridWidth = cols * tileSize;
            int totalGridHeight = rows * tileSize;

            int offsetX = (getWidth() - totalGridWidth) / 2;
            int offsetY = (getHeight() - totalGridHeight) / 2;

            Color backgroundBrown = new Color(155, 125, 90, 95);

            // 1. Draw underlying baseline network grid
            g2.setStroke(new BasicStroke(1.0f));
            g2.setColor(backgroundBrown);
            for (int x = 0; x < cols; x++) {
                for (int y = 0; y < rows; y++) {
                    int currentX = offsetX + (x * tileSize) + (tileSize / 2);
                    int currentY = offsetY + (y * tileSize) + (tileSize / 2);

                    if (x + 1 < cols) {
                        g2.drawLine(currentX, currentY, currentX + tileSize, currentY);
                    }
                    if (y + 1 < rows) {
                        g2.drawLine(currentX, currentY, currentX, currentY + tileSize);
                    }
                }
            }

            // 2. Render Active Purple Nodes and Pheromone Trails
            for (int x = 0; x < cols; x++) {
                for (int y = 0; y < rows; y++) {
                    double pheromone = grid.getPheromone(x, y);
                    int currentX = offsetX + (x * tileSize) + (tileSize / 2);
                    int currentY = offsetY + (y * tileSize) + (tileSize / 2);

                    boolean isCurrentNest = (x == controller.getNestX() && y == controller.getNestY());
                    boolean isCurrentFood = (x == controller.getFoodX() && y == controller.getFoodY());

                    if (pheromone > 0 || isCurrentNest || isCurrentFood) {
                        float intensity = (pheromone > 0) ? (float) Math.min(pheromone / 6.0, 1.0) : 0.5f;
                        int alphaVal = (int) (intensity * 200) + 55;
                        
                        if (pheromone > 0) {
                            g2.setColor(new Color(130, 0, 180, alphaVal));
                            int radius = (int) (tileSize * 0.35 + (intensity * tileSize * 0.35));
                            g2.fillOval(currentX - radius / 2, currentY - radius / 2, radius, radius);
                        }

                        int[][] directionalNeighbors = {{x + 1, y}, {x - 1, y}, {x, y + 1}, {x, y - 1}};
                        for (int[] n : directionalNeighbors) {
                            if (n[0] >= 0 && n[0] < cols && n[1] >= 0 && n[1] < rows) {
                                double nextPhero = grid.getPheromone(n[0], n[1]);
                                boolean isNeighborNest = (n[0] == controller.getNestX() && n[1] == controller.getNestY());
                                boolean isNeighborFood = (n[0] == controller.getFoodX() && n[1] == controller.getFoodY());

                                if ((pheromone > 0 && nextPhero > 0) || 
                                    (pheromone > 0 && (isNeighborNest || isNeighborFood)) || 
                                    (nextPhero > 0 && (isCurrentNest || isCurrentFood))) {
                                    
                                    float finalIntensity = Math.max(intensity, (float) Math.min(nextPhero / 6.0, 1.0));
                                    
                                    g2.setStroke(new BasicStroke(1.0f + (finalIntensity * 2.0f)));
                                    g2.setColor(new Color(130, 0, 180, (int) (finalIntensity * 135)));
                                    g2.drawLine(currentX, currentY, 
                                                offsetX + (n[0] * tileSize) + (tileSize / 2), 
                                                offsetY + (n[1] * tileSize) + (tileSize / 2));
                                }
                            }
                        }
                    }
                }
            }
            g2.setStroke(new BasicStroke(1.0f)); 

            // 3. Draw Nest Hub
            int nestX = offsetX + (controller.getNestX() * tileSize);
            int nestY = offsetY + (controller.getNestY() * tileSize);
            int elementMargin = Math.max(2, tileSize / 10);
            
            g2.setColor(new Color(140, 70, 20)); 
            g2.fillOval(nestX + elementMargin, nestY + elementMargin, tileSize - (elementMargin * 2), tileSize - (elementMargin * 2));
            g2.setColor(Color.WHITE);
            
            int fontSize = Math.max(9, tileSize / 3);
            g2.setFont(new Font("SansSerif", Font.BOLD, fontSize));
            FontMetrics metrics = g2.getFontMetrics();
            int stringX = nestX + (tileSize - metrics.stringWidth("N")) / 2;
            int stringY = nestY + ((tileSize - metrics.getHeight()) / 2) + metrics.getAscent();
            g2.drawString("N", stringX, stringY);

            // 4. Draw Food Anchor
            int foodX = offsetX + (controller.getFoodX() * tileSize);
            int foodY = offsetY + (controller.getFoodY() * tileSize);
            g2.setColor(Color.GREEN.darker());
            g2.fillRect(foodX + elementMargin, foodY + elementMargin, tileSize - (elementMargin * 2), tileSize - (elementMargin * 2));
            g2.setColor(Color.WHITE);
            stringX = foodX + (tileSize - metrics.stringWidth("F")) / 2;
            stringY = foodY + ((tileSize - metrics.getHeight()) / 2) + metrics.getAscent();
            g2.drawString("F", stringX, stringY);

            // 5. Draw Ants
            ArrayList<Ant> antList = controller.getAnts();
            for (int i = 0; i < antList.size(); i++) {
                Ant ant = antList.get(i);
                int size = Math.max(6, tileSize / 2);
                int drawX = offsetX + (ant.getX() * tileSize) + (tileSize / 2) - (size / 2);
                int drawY = offsetY + (ant.getY() * tileSize) + (tileSize / 2) - (size / 2);

                if (ant.hasFood()) {
                    g2.setColor(Color.GREEN); 
                } else {
                    g2.setColor(Color.BLACK); 
                }
                g2.fillOval(drawX, drawY, size, size);
            }
        }
    }

    public static void main(String[] args) {
        SwingUtilities.invokeLater(() -> new AntSimulationGUI());
    }
}