File size: 2,906 Bytes
80b67bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
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 |