vacuum_simulation/vacuum_simulation.py

import sys
import time
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QComboBox, QCheckBox, QLabel, QWidget, 
                             QTextEdit, QSplitter, QFrame)
from PyQt5.QtCore import QTimer, Qt
from PyQt5.QtGui import QColor, QPainter, QBrush, QPen

from environment import Environment, CellType, Direction
from search_algorithms import SearchAlgorithms

class GridWidget(QWidget):
    def __init__(self, rows, cols, cell_size, environment):
        super().__init__()
        self.rows = rows
        self.cols = cols
        self.cell_size = cell_size
        self.env = environment
        self.current_path = []
        self.current_direction = None
        
        self.setFixedSize(cols * cell_size, rows * cell_size)
    
    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        # Draw grid cells
        for row in range(self.rows):
            for col in range(self.cols):
                x = col * self.cell_size
                y = row * self.cell_size
                
                # Determine cell color based on type
                cell_type = self.env.grid[row][col]
                if cell_type == CellType.CLEAN:
                    color = QColor(255, 255, 0)  # Yellow
                elif cell_type == CellType.DIRTY:
                    color = QColor(255, 0, 0)    # Red
                elif cell_type == CellType.OBSTACLE:
                    color = QColor(0, 0, 255)    # Blue
                elif cell_type == CellType.EXPLORED:
                    color = QColor(0, 255, 0)    # Green
                
                # Draw cell
                painter.fillRect(x, y, self.cell_size, self.cell_size, color)
                painter.setPen(Qt.black)
                painter.drawRect(x, y, self.cell_size, self.cell_size)
        
        # Draw path
        if self.current_path:
            painter.setPen(QPen(QColor(255, 165, 0), 3))  # Orange, thicker line
            painter.setBrush(QBrush(QColor(255, 165, 0)))
            
            # Draw path lines
            for i in range(len(self.current_path) - 1):
                row1, col1 = self.current_path[i]
                row2, col2 = self.current_path[i + 1]
                
                x1 = col1 * self.cell_size + self.cell_size // 2
                y1 = row1 * self.cell_size + self.cell_size // 2
                x2 = col2 * self.cell_size + self.cell_size // 2
                y2 = row2 * self.cell_size + self.cell_size // 2
                
                painter.drawLine(x1, y1, x2, y2)
            
            # Draw path nodes (larger dots)
            for i, (row, col) in enumerate(self.current_path):
                x = col * self.cell_size + self.cell_size // 2
                y = row * self.cell_size + self.cell_size // 2
                
                # Make start and end points different
                if i == 0:  # Start
                    painter.setBrush(QBrush(QColor(0, 255, 0)))  # Green
                    radius = 6
                elif i == len(self.current_path) - 1:  # End
                    painter.setBrush(QBrush(QColor(255, 0, 0)))  # Red
                    radius = 6
                else:  # Intermediate
                    painter.setBrush(QBrush(QColor(255, 165, 0)))  # Orange
                    radius = 4
                
                painter.drawEllipse(x - radius, y - radius, radius * 2, radius * 2)
        
        # Draw vacuum
        if self.env.vacuum_pos:
            row, col = self.env.vacuum_pos
            x = col * self.cell_size
            y = row * self.cell_size
            
            # Draw vacuum as a circle with direction indicator
            painter.setPen(QPen(Qt.black, 2))
            painter.setBrush(QBrush(Qt.white))
            painter.drawEllipse(x + 5, y + 5, self.cell_size - 10, self.cell_size - 10)
            
            # Draw direction indicator
            if self.current_direction is not None:
                center_x = x + self.cell_size // 2
                center_y = y + self.cell_size // 2
                
                # Thicker direction indicator
                direction_pen = QPen(Qt.black, 3)
                painter.setPen(direction_pen)
                
                # Use the Direction enum properly
                if self.current_direction == Direction.UP:
                    painter.drawLine(center_x, center_y + 5, center_x, y + 5)
                elif self.current_direction == Direction.RIGHT:
                    painter.drawLine(center_x - 5, center_y, x + self.cell_size - 5, center_y)
                elif self.current_direction == Direction.DOWN:
                    painter.drawLine(center_x, center_y - 5, center_x, y + self.cell_size - 5)
                elif self.current_direction == Direction.LEFT:
                    painter.drawLine(center_x + 5, center_y, x + 5, center_y)
    
    def update_path(self, path, direction):
        self.current_path = path
        self.current_direction = direction
        self.update()

class VacuumSimulation(QMainWindow):
    def __init__(self, rows=15, cols=15):
        super().__init__()
        self.rows = rows
        self.cols = cols
        self.cell_size = 30
        self.env = Environment(rows, cols)
        self.search = SearchAlgorithms(self.env)
        
        # Simulation state
        self.current_path = []
        self.explored_cells = set()
        self.steps_taken = 0
        self.total_cost = 0
        self.current_direction = None
        self.is_running = False
        self.timer = QTimer()
        self.timer.timeout.connect(self.next_step)
        
        # Performance metrics
        self.total_nodes_expanded = 0
        self.total_computation_time = 0
        self.algorithm_stats = {
            "BFS": {"runs": 0, "total_nodes": 0, "total_time": 0},
            "A* Manhattan": {"runs": 0, "total_nodes": 0, "total_time": 0},
            "A* Euclidean": {"runs": 0, "total_nodes": 0, "total_time": 0},
            "A* Chebyshev": {"runs": 0, "total_nodes": 0, "total_time": 0}
        }
        
        self.init_ui()
        self.update_display()
    
    def init_ui(self):
        self.setWindowTitle("Vacuum Cleaner Search Simulation - Algorithm Comparison")
        
        # Use a larger window to accommodate side panel
        grid_width = self.cols * self.cell_size
        grid_height = self.rows * self.cell_size
        self.setMinimumSize(grid_width + 400, grid_height + 200)
        
        # Central widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # Main layout using splitter for resizable panels
        main_layout = QHBoxLayout()
        central_widget.setLayout(main_layout)
        
        # Left side: Grid and controls
        left_widget = QWidget()
        left_layout = QVBoxLayout()
        left_widget.setLayout(left_layout)
        
        # Right side: Information panel
        right_widget = QWidget()
        right_widget.setMaximumWidth(350)
        right_layout = QVBoxLayout()
        right_widget.setLayout(right_layout)
        
        # === LEFT PANEL: Grid and Controls ===
        
        # Metrics display at top
        metrics_layout = QHBoxLayout()
        
        # Left metrics column
        left_metrics = QVBoxLayout()
        self.steps_label = QLabel("Steps: 0")
        self.cost_label = QLabel("Total Cost: 0.0")
        self.turn_cost_label = QLabel("Turn Cost: 0.0")
        self.dirty_label = QLabel("Dirty Cells: 0")
        
        left_metrics.addWidget(self.steps_label)
        left_metrics.addWidget(self.cost_label)
        left_metrics.addWidget(self.turn_cost_label)
        left_metrics.addWidget(self.dirty_label)
        
        # Right metrics column
        right_metrics = QVBoxLayout()
        self.nodes_explored_label = QLabel("Nodes Explored: 0")
        self.nodes_expanded_label = QLabel("Nodes Expanded: 0")
        self.comp_time_label = QLabel("Comp Time: 0.000s")
        self.algorithm_label = QLabel("Algorithm: None")
        
        right_metrics.addWidget(self.nodes_explored_label)
        right_metrics.addWidget(self.nodes_expanded_label)
        right_metrics.addWidget(self.comp_time_label)
        right_metrics.addWidget(self.algorithm_label)
        
        metrics_layout.addLayout(left_metrics)
        metrics_layout.addLayout(right_metrics)
        metrics_layout.addStretch()
        
        left_layout.addLayout(metrics_layout)
        
        # Control panel
        control_layout = QHBoxLayout()
        
        # Reset button
        self.reset_button = QPushButton("Reset")
        self.reset_button.clicked.connect(self.reset_simulation)
        control_layout.addWidget(self.reset_button)
        
        # Next button
        self.next_button = QPushButton("Next")
        self.next_button.clicked.connect(self.next_step)
        control_layout.addWidget(self.next_button)
        
        # Run button
        self.run_button = QPushButton("Run")
        self.run_button.clicked.connect(self.toggle_run)
        control_layout.addWidget(self.run_button)
        
        control_layout.addStretch()
        left_layout.addLayout(control_layout)
        
        # Search options
        search_layout = QHBoxLayout()
        
        # Turn cost checkbox
        self.turn_cost_checkbox = QCheckBox("Turn Cost (0.5 per 90° turn)")
        self.turn_cost_checkbox.stateChanged.connect(self.toggle_turn_cost)
        search_layout.addWidget(self.turn_cost_checkbox)
        
        # Search algorithm dropdown
        search_layout.addWidget(QLabel("Search:"))
        self.search_combo = QComboBox()
        self.search_combo.addItems(["BFS", "A* Manhattan", "A* Euclidean", "A* Chebyshev"])
        search_layout.addWidget(self.search_combo)
        
        search_layout.addStretch()
        left_layout.addLayout(search_layout)
        
        # Grid display
        self.grid_widget = GridWidget(self.rows, self.cols, self.cell_size, self.env)
        left_layout.addWidget(self.grid_widget)
        
        # Legend (moved to bottom of left panel)
        legend_layout = QHBoxLayout()
        
        def create_legend_item(color, text):
            item_widget = QWidget()
            item_layout = QHBoxLayout()
            item_widget.setLayout(item_layout)
            
            color_label = QLabel()
            color_label.setFixedSize(16, 16)
            color_label.setStyleSheet(f"background-color: {color}; border: 1px solid black")
            item_layout.addWidget(color_label)
            
            text_label = QLabel(text)
            text_label.setStyleSheet("font-size: 10px;")
            item_layout.addWidget(text_label)
            item_layout.setContentsMargins(2, 0, 5, 0)
            
            return item_widget
        
        legend_layout.addWidget(create_legend_item("yellow", "Clean"))
        legend_layout.addWidget(create_legend_item("red", "Dirty"))
        legend_layout.addWidget(create_legend_item("blue", "Obstacle"))
        legend_layout.addWidget(create_legend_item("green", "Explored"))
        legend_layout.addWidget(create_legend_item("orange", "Path"))
        legend_layout.addStretch()
        
        left_layout.addLayout(legend_layout)
        
        # === RIGHT PANEL: Information and Statistics ===
        
        # Algorithm Information
        info_label = QLabel("Algorithm Information:")
        info_label.setStyleSheet("font-weight: bold; font-size: 12px; margin-bottom: 5px;")
        right_layout.addWidget(info_label)
        
        info_text = QTextEdit()
        info_text.setMaximumHeight(120)
        info_text.setReadOnly(True)
        info_text.setHtml("""
        <b>Search Algorithms:</b><br>
        • <b>BFS</b>: Explores all directions equally, finds shortest path<br>
        • <b>A* Manhattan</b>: Uses city-block distance heuristic<br>
        • <b>A* Euclidean</b>: Uses straight-line distance heuristic<br>
        • <b>A* Chebyshev</b>: Uses chessboard distance heuristic<br><br>
        <b>Turn Cost</b>: When enabled, 90° turns cost +0.5
        """)
        right_layout.addWidget(info_text)
        
        # Performance Statistics
        stats_label = QLabel("Algorithm Performance:")
        stats_label.setStyleSheet("font-weight: bold; font-size: 12px; margin-top: 10px; margin-bottom: 5px;")
        right_layout.addWidget(stats_label)
        
        self.stats_text = QTextEdit()
        self.stats_text.setReadOnly(True)
        self.stats_text.setStyleSheet("font-family: monospace; font-size: 10px;")
        right_layout.addWidget(self.stats_text)
        
        # Current Run Analysis
        analysis_label = QLabel("Current Run Analysis:")
        analysis_label.setStyleSheet("font-weight: bold; font-size: 12px; margin-top: 10px; margin-bottom: 5px;")
        right_layout.addWidget(analysis_label)
        
        self.analysis_text = QTextEdit()
        self.analysis_text.setMaximumHeight(100)
        self.analysis_text.setReadOnly(True)
        self.analysis_text.setStyleSheet("font-family: monospace; font-size: 10px;")
        right_layout.addWidget(self.analysis_text)
        
        # Add stretch to push everything to top
        right_layout.addStretch()
        
        # Add both panels to main layout
        main_layout.addWidget(left_widget)
        main_layout.addWidget(right_widget)
        
        # Set left widget to expand, right widget fixed width
        main_layout.setStretchFactor(left_widget, 1)
        main_layout.setStretchFactor(right_widget, 0)
        
        self.update_stats_display()
    
    def update_stats_display(self):
        stats_text = "<pre>"
        for algo, stats in self.algorithm_stats.items():
            if stats["runs"] > 0:
                avg_nodes = stats["total_nodes"] / stats["runs"]
                avg_time = stats["total_time"] / stats["runs"]
                stats_text += f"{algo:<15} {stats['runs']:>3} runs\n"
                stats_text += f"  Avg Nodes: {avg_nodes:>6.1f}\n"
                stats_text += f"  Avg Time:  {avg_time:>7.4f}s\n"
            else:
                stats_text += f"{algo:<15} No runs yet\n"
        stats_text += "</pre>"
        
        # Add Chebyshev vs Euclidean comparison to analysis
        chebyshev_stats = self.algorithm_stats["A* Chebyshev"]
        euclidean_stats = self.algorithm_stats["A* Euclidean"]
        
        analysis_text = "<pre>"
        if chebyshev_stats["runs"] > 0 and euclidean_stats["runs"] > 0:
            chebyshev_avg_nodes = chebyshev_stats["total_nodes"] / chebyshev_stats["runs"]
            euclidean_avg_nodes = euclidean_stats["total_nodes"] / euclidean_stats["runs"]
            chebyshev_avg_time = chebyshev_stats["total_time"] / chebyshev_stats["runs"]
            euclidean_avg_time = euclidean_stats["total_time"] / euclidean_stats["runs"]
            
            if euclidean_avg_nodes > 0:
                node_ratio = chebyshev_avg_nodes / euclidean_avg_nodes
                analysis_text += f"Node Exploration:\n"
                analysis_text += f"  Chebyshev explores {node_ratio:.1f}x\n"
                analysis_text += f"  more nodes than Euclidean\n\n"
            
            if euclidean_avg_time > 0:
                time_ratio = chebyshev_avg_time / euclidean_avg_time
                analysis_text += f"Computation Time:\n"
                analysis_text += f"  Chebyshev is {time_ratio:.1f}x\n"
                analysis_text += f"  slower than Euclidean"
        else:
            analysis_text += "Run simulations to see\nalgorithm comparisons"
        analysis_text += "</pre>"
        
        self.stats_text.setHtml(stats_text)
        self.analysis_text.setHtml(analysis_text)
    
    def update_display(self):
        self.steps_label.setText(f"Steps: {self.steps_taken}")
        self.cost_label.setText(f"Total Cost: {self.total_cost:.2f}")
        
        # Calculate turn cost separately
        turn_cost = max(0, self.total_cost - self.steps_taken)
        self.turn_cost_label.setText(f"Turn Cost: {turn_cost:.2f}")
        
        self.dirty_label.setText(f"Dirty Cells: {self.env.get_dirty_count()}")
        self.nodes_explored_label.setText(f"Nodes Explored: {len(self.explored_cells)}")
        self.nodes_expanded_label.setText(f"Nodes Expanded: {self.total_nodes_expanded}")
        self.comp_time_label.setText(f"Comp Time: {self.total_computation_time:.3f}s")
        
        current_algorithm = self.search_combo.currentText()
        self.algorithm_label.setText(f"Algorithm: {current_algorithm}")
        
        self.grid_widget.update_path(self.current_path, self.current_direction)
    
    def reset_simulation(self):
        self.env.reset()
        self.current_path = []
        self.explored_cells = set()
        self.steps_taken = 0
        self.total_cost = 0
        self.current_direction = None
        self.total_nodes_expanded = 0
        self.total_computation_time = 0
        self.is_running = False
        self.timer.stop()
        self.run_button.setText("Run")
        self.update_display()
    
    def next_step(self):
        if self.env.is_clean():
            self.is_running = False
            self.timer.stop()
            self.run_button.setText("Run")
            return
        
        # If we don't have a path, find one
        if not self.current_path:
            self.find_path()
            # If still no path after searching, stop
            if not self.current_path:
                self.is_running = False
                self.timer.stop()
                self.run_button.setText("Run")
                return
        
        # If we have a path, move along it
        if self.current_path:
            # Move to next position in path
            next_pos = self.current_path.pop(0)
            current_pos = self.env.vacuum_pos
            
            # Calculate movement cost with turn cost
            move_cost = 1  # Base movement cost
            
            # Determine new direction based on movement
            new_direction = Direction.from_movement(current_pos, next_pos)
            
            if self.turn_cost_checkbox.isChecked() and self.current_direction is not None:
                # Add turn cost if direction changed
                if self.current_direction != new_direction:
                    turn_cost = 0.5  # 90° turn cost
                    move_cost += turn_cost
            
            # Update direction
            self.current_direction = new_direction
            
            # Update vacuum position
            self.env.vacuum_pos = next_pos
            self.steps_taken += 1
            self.total_cost += move_cost
            
            # Clean the cell if it's dirty
            if self.env.clean_cell(next_pos[0], next_pos[1]):
                # If we cleaned a cell, we need to find a new path
                self.current_path = []
            
            # Mark cell as explored
            self.env.mark_explored(next_pos[0], next_pos[1])
        
        self.update_display()
    
    def find_path(self):
        # Get selected search algorithm
        search_type = self.search_combo.currentText()
        
        if search_type == "BFS":
            algorithm = "bfs"
            heuristic = "manhattan"
        else:
            algorithm = "a_star"
            if search_type == "A* Manhattan":
                heuristic = "manhattan"
            elif search_type == "A* Euclidean":
                heuristic = "euclidean"
            else:  # A* Chebyshev
                heuristic = "chebyshev"
        
        # Update search algorithm with current turn cost setting
        self.search.turn_cost_enabled = self.turn_cost_checkbox.isChecked()
        
        # Find path to nearest dirty cell
        path, explored, nodes_expanded, computation_time = self.search.find_path_to_nearest_dirty(
            self.env.vacuum_pos, algorithm, heuristic)
        
        # Update performance metrics
        self.total_nodes_expanded += nodes_expanded
        self.total_computation_time += computation_time
        
        # Update algorithm statistics
        self.algorithm_stats[search_type]["runs"] += 1
        self.algorithm_stats[search_type]["total_nodes"] += nodes_expanded
        self.algorithm_stats[search_type]["total_time"] += computation_time
        
        if path:
            # Remove current position from path
            self.current_path = path[1:]
            self.explored_cells.update(explored)
            
            # Mark explored cells
            for row, col in explored:
                if (row, col) != self.env.vacuum_pos and self.env.grid[row][col] == CellType.CLEAN:
                    self.env.grid[row][col] = CellType.EXPLORED
        
        self.update_stats_display()
    
    def toggle_run(self):
        self.is_running = not self.is_running
        
        if self.is_running:
            self.run_button.setText("Pause")
            self.timer.start(500)  # 0.5 second interval for faster visualization
        else:
            self.run_button.setText("Run")
            self.timer.stop()
    
    def toggle_turn_cost(self):
        # When turn cost is toggled, we need to recalculate the path
        if self.current_path:
            self.current_path = []