vacuum_simulation/V0/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

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
        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