SI_2020/main.py

import pygame
import warehouse
import agent
import random
import sys
from attributes import PackSize, PackStatus, COLORS

WINDOW_SIZE = (600, 600)
COLOR_OF_FIELD = {
            'Floor': 'gray',
            'Rack': 'white',
            'Pack': 'yellow'
        }
TILE_WIDTH = 30
TILE_HEIGHT = 30
CIRCLE_CENTER_X, CIRCLE_CENTER_Y = int(TILE_WIDTH/2), int(TILE_HEIGHT/2)

class MainGameFrame:
    def __init__(self):
        self.display = pygame.display.set_mode(WINDOW_SIZE)
        agent_radius = int(TILE_WIDTH/2)
        self.warehouse_map = warehouse.Warehouse(20, 20, 150, 20)
        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.agent.move()
            pygame.display.update()
            self.clock.tick(5)

    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]
        if not isinstance(current_tile, warehouse.Tile):
            current_tile = current_tile.lays_on_field if isinstance(current_tile, warehouse.Pack) else None
        color = COLOR_OF_FIELD.get(current_tile.category.name, 'white')
        color = COLORS[color]
        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))

    def draw_agent(self):
        agent_position_x, agent_position_y = self.agent.x, self.agent.y
        agent_screen_position = ((agent_position_x*TILE_WIDTH) + CIRCLE_CENTER_X, (agent_position_y*TILE_HEIGHT) + CIRCLE_CENTER_Y)
        pygame.draw.circle(self.display, COLORS['black'], agent_screen_position, self.agent.radius, int(self.agent.radius/2))

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