import pygame
from treelearn import treelearn
import loadmodel
from astar import astar
from state import State
import time
from garbage_truck import GarbageTruck
from heuristicfn import heuristicfn
from map import randomize_map
from tree import apply_tree
from genetic import apply_genetic


pygame.init()
WIDTH, HEIGHT = 800, 800
window = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Intelligent Garbage Collector")
AGENT_IMG = pygame.image.load("garbage-truck-nbg.png")
AGENT = pygame.transform.scale(AGENT_IMG, (50, 50))
FPS = 10
FIELDCOUNT = 16
FIELDWIDTH = 50
BASE_IMG = pygame.image.load("Tiles/Base.jpg")
BASE = pygame.transform.scale(BASE_IMG, (50, 50))

def draw_window(agent, fields, flip, turn):
    if flip:
        direction = pygame.transform.flip(AGENT, True, False)
        if turn:
            direction = pygame.transform.rotate(AGENT, -90)
    else:
        direction = pygame.transform.flip(AGENT, False, False)
        if turn:
            direction = pygame.transform.rotate(AGENT, 90)
    for i in range(16):
        for j in range(16):
            window.blit(fields[i][j], (i * 50, j * 50))
    window.blit(direction, (agent.rect.x, agent.rect.y))  # wyswietlanie agenta
    pygame.display.update()


def main():
    clf = treelearn()
    clock = pygame.time.Clock()
    run = True
    fields, priority_array, request_list, imgpath_array = randomize_map()
    apply_tree(request_list)
    apply_genetic(request_list)
    agent = GarbageTruck(0, 0, pygame.Rect(0, 0, 50, 50), 0, request_list, clf)   # tworzenie pola dla agenta
    low_space = 0
    while run:
        clock.tick(FPS)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
        draw_window(agent, fields, False, False)   # false = kierunek east (domyslny), true = west
        x, y = agent.next_destination()
        if x == agent.rect.x and y == agent.rect.y:
            print('out of jobs')
            break
        if low_space == 1:
            x, y = 0, 0
        steps = astar(State(None, None, agent.rect.x, agent.rect.y, agent.orientation,
                            priority_array[agent.rect.x//50][agent.rect.y//50],
                            heuristicfn(agent.rect.x, agent.rect.y, x, y)), x, y, priority_array)
        for interm in steps:
            if interm.action == 'LEFT':
                agent.turn_left()
                if agent.orientation == 0:
                    draw_window(agent, fields, False, False)
                elif agent.orientation == 2:
                    draw_window(agent, fields, True, False)
                elif agent.orientation == 1:
                    draw_window(agent, fields, True, True)
                else:
                    draw_window(agent, fields, False, True)
            elif interm.action == 'RIGHT':
                agent.turn_right()
                if agent.orientation == 0:
                    draw_window(agent, fields, False, False)
                elif agent.orientation == 2:
                    draw_window(agent, fields, True, False)
                elif agent.orientation == 1:
                    draw_window(agent, fields, True, True)
                else:
                    draw_window(agent, fields, False, True)
            elif interm.action == 'FORWARD':
                agent.forward()
                if agent.orientation == 0:
                    draw_window(agent, fields, False, False)
                elif agent.orientation == 2:
                    draw_window(agent, fields, True, False)
                elif agent.orientation == 1:
                    draw_window(agent, fields, True, True)
                else:
                    draw_window(agent, fields, False, True)
            time.sleep(0.3)
        if (agent.rect.x // 50 != 0) or (agent.rect.y // 50 != 0):
            garbage_type = loadmodel.classify(imgpath_array[agent.rect.x // 50][agent.rect.y // 50])
            low_space = agent.collect(garbage_type)

        fields[agent.rect.x//50][agent.rect.y//50] = BASE
        priority_array[agent.rect.x//50][agent.rect.y//50] = 100
        time.sleep(0.5)

    pygame.quit()


if __name__ == "__main__":
    main()