import pygame import warehouse import agent import random import sys from attributes import PackStatus, COLORS, DIRECTION_ANGLES WINDOW_SIZE = (640, 640) COLOR_OF_FIELD = { 'Floor': 'gray', 'Rack': 'white', 'Pack': 'yellow', 'path': 'orange', 'FridgeFloor': 'lightblue', 'Fridge': 'iceblue' } TILE_WIDTH = 32 TILE_HEIGHT = 32 CIRCLE_CENTER_X, CIRCLE_CENTER_Y = int(TILE_WIDTH/2), int(TILE_HEIGHT/2) class MainGameFrame: def __init__(self): pygame.font.init() self.display = pygame.display.set_mode(WINDOW_SIZE) pygame.display.set_caption("Smart ForkLift") agent_radius = int(TILE_WIDTH/2) self.agent_tex = pygame.image.load('forklift.png') self.font = pygame.font.Font('freesansbold.ttf', 16) self.warehouse_map = warehouse.Warehouse(20, 20, 150, 10) starting_x, starting_y = self.set_starting_agent_position() self.agent = agent.Agent(starting_x, starting_y, self.warehouse_map, agent_radius) self.clock = pygame.time.Clock() def run(self): while True: for event in pygame.event.get(): if event.type == pygame.QUIT: pygame.quit() sys.exit() self.draw_floor() self.draw_packages() self.draw_agent() self.draw_nums() self.agent.move() packages_to_go = [p for p in self.warehouse_map.packages if p.status != PackStatus.STORED] if not packages_to_go and not self.agent.transported_package: pygame.quit() sys.exit() pygame.display.update() self.clock.tick(6) def draw_floor(self): for x in range(self.warehouse_map.width): for y in range(self.warehouse_map.height): self.draw_field(x, y) def draw_field(self, x, y): current_tile = self.warehouse_map.tiles[x][y] color = COLOR_OF_FIELD.get(current_tile.category.name, 'white') color = COLORS[color] if (current_tile.x_position,current_tile.y_position) in [(a.x, a.y) for a in self.agent.path]: color = COLORS.get('orange') pygame.draw.rect(self.display, COLORS['black'], (x * TILE_WIDTH, y * TILE_HEIGHT, TILE_WIDTH, TILE_HEIGHT)) pygame.draw.rect(self.display, color, ((x * TILE_WIDTH) + 1, (y * TILE_HEIGHT) + 1, TILE_WIDTH - 1, TILE_HEIGHT - 1)) def draw_packages(self): def get_package_color(pack): colors = { PackStatus.LOOSE: COLORS['yellow'], PackStatus.STORED: COLORS['lightgreen'], PackStatus.STORED_BAD_LOCATION: COLORS['red'] } return colors[pack.status] for pack in self.warehouse_map.packages: pack_x, pack_y = pack.lays_on_field.x_position, pack.lays_on_field.y_position package_color = get_package_color(pack) pygame.draw.rect(self.display, package_color, ((pack_x * TILE_WIDTH) + 3, (pack_y * TILE_HEIGHT) + 3, TILE_WIDTH - 5, TILE_HEIGHT - 5)) if pack.category == "freezed": pygame.draw.rect(self.display, COLORS['blue'], ((pack_x * TILE_WIDTH) + 2, (pack_y * TILE_HEIGHT) + 2, TILE_WIDTH - 4, TILE_HEIGHT - 4), 3) def draw_nums(self): for row in self.warehouse_map.tiles: for cell in row: if cell.category.name in self.warehouse_map.storage_types: text_surface = self.font.render(str(cell.capacity), True, (0, 0, 0)) self.display.blit(text_surface, ((cell.x_position * TILE_WIDTH) + 6, (cell.y_position * TILE_HEIGHT) + 6)) def draw_agent(self): rotated = pygame.transform.rotate(self.agent_tex, DIRECTION_ANGLES.get(self.agent.direction)) self.display.blit(rotated, (self.agent.x*TILE_WIDTH, self.agent.y*TILE_WIDTH)) def set_starting_agent_position(self): starting_x, starting_y = random.randrange(self.warehouse_map.width), random.randrange(self.warehouse_map.height) while not isinstance(self.warehouse_map.tiles[starting_x][starting_y], warehouse.Tile) or self.warehouse_map.tiles[starting_x][starting_y].category.name != 'Floor': starting_x, starting_y = random.randrange(self.warehouse_map.width), random.randrange( self.warehouse_map.height) return starting_x, starting_y if __name__ == '__main__': maingame = MainGameFrame() maingame.run()