vacuum_simulation/environment.py

import random
from enum import Enum

class CellType(Enum):
    CLEAN = 0
    DIRTY = 1
    OBSTACLE = 2
    EXPLORED = 3

class Direction(Enum):
    UP = 0
    RIGHT = 1
    DOWN = 2
    LEFT = 3
    
    @classmethod
    def from_movement(cls, current_pos, next_pos):
        """Determine direction from current position to next position"""
        current_row, current_col = current_pos
        next_row, next_col = next_pos
        
        if next_row < current_row:
            return cls.UP
        elif next_row > current_row:
            return cls.DOWN
        elif next_col > current_col:
            return cls.RIGHT
        elif next_col < current_col:
            return cls.LEFT
        return None

class Environment:
    def __init__(self, rows, cols, obstacle_density=0.2, dirt_density=0.3):
        self.rows = rows
        self.cols = cols
        self.grid = [[CellType.CLEAN for _ in range(cols)] for _ in range(rows)]
        self.obstacle_density = obstacle_density
        self.dirt_density = dirt_density
        self.vacuum_pos = None
        self.dirty_cells = set()
        
        self.generate_environment()
    
    def generate_environment(self):
        # Place obstacles
        for i in range(self.rows):
            for j in range(self.cols):
                if random.random() < self.obstacle_density:
                    self.grid[i][j] = CellType.OBSTACLE
        
        # Place dirt on clean cells only
        for i in range(self.rows):
            for j in range(self.cols):
                if self.grid[i][j] == CellType.CLEAN and random.random() < self.dirt_density:
                    self.grid[i][j] = CellType.DIRTY
                    self.dirty_cells.add((i, j))
        
        # Place vacuum at a random clean position
        clean_positions = [(i, j) for i in range(self.rows) for j in range(self.cols) 
                          if self.grid[i][j] == CellType.CLEAN]
        if clean_positions:
            self.vacuum_pos = random.choice(clean_positions)
    
    def reset(self):
        self.grid = [[CellType.CLEAN for _ in range(self.cols)] for _ in range(self.rows)]
        self.dirty_cells = set()
        self.generate_environment()
    
    def is_valid_position(self, row, col):
        return (0 <= row < self.rows and 
                0 <= col < self.cols and 
                self.grid[row][col] != CellType.OBSTACLE)
    
    def clean_cell(self, row, col):
        if (row, col) in self.dirty_cells:
            self.grid[row][col] = CellType.CLEAN
            self.dirty_cells.remove((row, col))
            return True
        return False
    
    def mark_explored(self, row, col):
        if self.grid[row][col] == CellType.CLEAN:
            self.grid[row][col] = CellType.EXPLORED
    
    def get_dirty_count(self):
        return len(self.dirty_cells)
    
    def is_clean(self):
        return len(self.dirty_cells) == 0