Machine_learning_2023/domain/world.py

from decisionTree.evaluate import evaluate
from domain.entities.entity import Entity


class World:
    def __init__(self, width: int, height: int) -> object:
        self.costs = [[1 for j in range(height)] for i in range(width)]
        self.width = width
        self.height = height
        self.dust = [[[] for j in range(height)] for i in range(width)]
        self.dustList = []
        self.obstacles = [[[] for j in range(height)] for i in range(width)]
        self.entity = [[[] for j in range(height)] for i in range(width)]

        self.vacuum = None
        self.cat = None
        self.doc_station = None

    def add_entity(self, entity: Entity):
        if entity.type == "DOC_STATION":
            self.doc_station = entity
        elif entity.type == "PEEL":
            self.dust[entity.x][entity.y].append(entity)
            self.dustList.append(Entity(entity.x, entity.y, "PEEL"))
        elif entity.type == "EARRING":
            self.dust[entity.x][entity.y].append(entity)
            self.dustList.append(Entity(entity.x, entity.y, "EARRING"))
        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)

        self.entity[entity.x][entity.y].append(entity)

    def is_entity_at(
        self,
        x: int,
        y: int,
    ) -> bool:
        if len(self.entity[x][y]) > 0:
            return True
        return False

    def delete_entities_at_Of_type(self, x: int, y: int, type: str):
        entities = self.entity[x][y]
        for entity in entities:
            if entity.type == type:
                entities.remove(entity)

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

    def garbage_at(self, x: int, y: int) -> list[Entity]:
        if len(self.dust[x][y]) == 0:
            return []
        return [i for i in self.dust[x][y] if evaluate([i.properties])[0] == 1]

    def is_docking_station_at(self, x: int, y: int) -> bool:
        return bool(self.doc_station.x == x and self.doc_station.y == y)

    def accepted_move(self, checking_x, checking_y):
        if (
            checking_x > self.width - 1
            or checking_y > self.height - 1
            or checking_x < 0
            or checking_y < 0
        ):
            return False

        if self.is_obstacle_at(checking_x, checking_y):
            return False

        return True

    def get_cost(self, x, y) -> float:
        return self.costs[x][y]