import pygame
import warehouse
import agent
from genetic_algorithm import gen_alg
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, mutation_prob=0.03, generation_size=40, number_of_generations=30, amount_of_promotion=0):
        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.agent_tex2 = pygame.image.load('forklift_loaded.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()
        packs_coords = [(pack.lays_on_field.x_position, pack.lays_on_field.y_position) for pack in self.warehouse_map.packages]
        list_of_racks = self.warehouse_map.get_all_racks(True)
        racks_coords = [(line.x_position, line.y_position) for line in list_of_racks]
        packs_sizes = [pack.size for pack in self.warehouse_map.packages]
        racks_capacities = [rack.capacity for rack in list_of_racks]
        # print("koordynaty paczek: ",packs_coords)
        # print("koordynaty regałów: ",racks_coords)
        print("wagi paczek: ",packs_sizes)
        print("pojemności regałów: ",racks_capacities)
        gen_alg(packs_sizes, racks_capacities, number_of_generations, generation_size, mutation_prob, amount_of_promotion, self.agent.location_classifier)
        self.agent.create_graph_map()
        self.agent.trace_route()
        self.agent.find_path()

    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))
        for package in self.warehouse_map.packages:
            if package.status == PackStatus.LOOSE:
                text_surface = self.font.render(str(package.id + 1), True, (0, 0, 0))
                self.display.blit(text_surface, ((package.lays_on_field.x_position * TILE_WIDTH) + 6, (package.lays_on_field.y_position * TILE_HEIGHT) + 6))

    def draw_agent(self):
        if self.agent.is_loaded:
            rotated = pygame.transform.rotate(self.agent_tex2, DIRECTION_ANGLES.get(self.agent.direction))
        else:
            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()