import pygame
from treelearn import treelearn


from astar import astar
from state import State
import time
from garbage_truck import GarbageTruck
from heuristicfn import heuristicfn
from map import randomize_map

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

GRASS_IMG = pygame.image.load("grass.png")
GRASS = pygame.transform.scale(GRASS_IMG, (50, 50))
def draw_window(agent, fields, flip):
    if flip:
        direction = pygame.transform.flip(AGENT, True, False)
    else:
        direction = pygame.transform.flip(AGENT, False, False)
    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 = randomize_map()
    agent = GarbageTruck(0, 0, pygame.Rect(0, 0, 50, 50), 0, request_list, clf)   # tworzenie pola dla agenta
    while run:
        clock.tick(FPS)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
        draw_window(agent, fields, 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
        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()
                draw_window(agent, fields, True)
            elif interm.action == 'RIGHT':
                agent.turn_right()
                draw_window(agent, fields, False)
            elif interm.action == 'FORWARD':
                agent.forward()
                if agent.orientation == 0:
                    draw_window(agent, fields, False)
                elif agent.orientation == 2:
                    draw_window(agent, fields, True)
                else:
                    draw_window(agent, fields, False)
            time.sleep(0.3)
        agent.collect()
        fields[agent.rect.x//50][agent.rect.y//50] = GRASS
        priority_array[agent.rect.x//50][agent.rect.y//50] = 1
        time.sleep(0.5)
    

    pygame.quit()


if __name__ == "__main__":
    main()