Machine_learning_2023/domain/world.py

from domain.entities.entity import Entity


class World:
    def __init__(self, width: int, height: int):
        self.width = width
        self.height = height
        self.dust = [
            [[] for j in range(height)] for i in range(width)
        ]
        self.obstacles = [
            [[] for j in range(height)] for i in range(width)
        ]

        self.vacuum = None
        self.cat = None

    def add_entity(self, entity: Entity):
        if entity.type == "PEEL":
            self.dust[entity.x][entity.y].append(entity)
        elif entity.type == "VACUUM":
            self.vacuum = entity
        elif entity.type == "CAT":
            self.cat = entity
            self.obstacles[entity.x][entity.y].append(entity)
        else:
            self.obstacles[entity.x][entity.y].append(entity)

    def is_obstacle_at(self, x: int, y: int) -> bool:
        return bool(self.obstacles[x][y])