from keras.models import Sequential, load_model
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.layers.normalization import BatchNormalization
from PIL import Image
from random import shuffle, choice
import numpy as np
import os

import tree
import pygame
from time import sleep
from os import path
from colors import gray
from house import create_houses
from truck import Truck
from trash import Trash
from TSP import tsp, tspmove
from bfs import bfs, distance
from pool import create_pools
from granny import create_granny

model = load_model("model.h5")

multi_trash = []
pygame.init()


def game_keys(truck, multi_trash, houses, auto=False):
    for event in pygame.event.get():
        if event.type == pygame.KEYDOWN:
            if auto:
                sleep(.01)
            if (event.key == pygame.K_w or event.key == pygame.K_UP):
                if truck.test_crash(houses):
                    break
                truck.move()
                print('↑')
                for tindex, trash in enumerate(multi_trash):
                    if truck.pos == trash.pos:
                        truck.mass += trash.mass
                        truck.space += trash.space
                        print(truck.mass, truck.space)
                        prediction = model.predict(multi_trash[tindex].content)
                        for i in range (3):
                            if multi_trash[tindex].names[i][:3] == 'cat':
                                truck.allCats = truck.allCats + 1
                            else:
                                truck.allTrash = truck.allTrash + 1
                            if prediction[i][0] > prediction[i][1]:
                                truck.container.append(multi_trash[tindex].names[i])
                                if multi_trash[tindex].names[i][:3] == 'cat':
                                    truck.cats = truck.cats + 1
                                else:
                                    truck.trash = truck.trash + 1
                        print(f'cats in truck: {truck.cats} cats encountered: {truck.allCats}')
                        print(f'trash in truck: {truck.trash} trash encountered: {truck.allTrash}')
                        print(truck.container)
                        multi_trash.pop(tindex)
                break
            elif (event.key == pygame.K_a or event.key == pygame.K_LEFT):
                print('←')
                truck.rotate(-1)
                truck.rotate_image(90)
                break
            elif (event.key == pygame.K_d or event.key == pygame.K_RIGHT):
                print('→')
                truck.rotate(1)
                truck.rotate_image(-90)
                break


def update_images(gameDisplay, truck, multi_trash, houses,pools,grannies):
    gameDisplay.fill(gray)
    for house in houses:
        gameDisplay.blit(pygame.image.load('./img/house.png'), house.pos)
    for pool in pools:
        gameDisplay.blit(pygame.image.load('./img/wet.png'), pool.pos)
    for granny in grannies:
        gameDisplay.blit(pygame.image.load('./img/granny.png'), granny.pos)
    for trash in multi_trash:
        gameDisplay.blit(pygame.image.load('./img/trash.png'), trash.pos)
    gameDisplay.blit(truck.image, truck.pos)
    pygame.display.update()


def game_loop():
    game_w = 1280  # 32
    game_h = 640  # 16
    grid_size = 40

    gameDisplay = pygame.display.set_mode((game_w, game_h))
    pygame.display.set_caption('Garbage truck')
    clock = pygame.time.Clock()

    gameExit = False

    truck = Truck(game_w, game_h, grid_size)
    houses = create_houses(grid_size)
    pools = create_pools(grid_size)
    grannies = create_granny(grid_size)

    for _ in range(7):
        trash = Trash(game_w, game_h, grid_size)
        trash.new_pos(truck.pos, houses, multi_trash)
        multi_trash.append(trash)

    model.predict(multi_trash[0].content)
    while not gameExit:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    quit()
                if (event.key == pygame.K_l):
                    if path.isfile('./tree_model') and not os.stat(
                            './tree_model').st_size == 0:
                        decision_tree = tree.load_tree_from_structure('./tree_model')
                        print("Tree model loaded!")

                if (event.key == pygame.K_k):
                    print(":>")
                    trash = multi_trash[0]
                    dis = [120,120]
                    dis_dump = distance(truck.pos,dis)
                    dis_trash = distance(truck.pos, trash.pos)
                    decision =  tree.making_decision(decision_tree,
                                                      dis_dump // 12 + 1,
                                                      dis_trash // 12 + 1,
                                                      truck.mass // 20 + 1, truck.space // 20 + 1,
                                                      trash.mass // 20 + 1,
                                                      trash.space // 20 + 1)

                    print(decision)
                    if(decision[0]==0):
                        actions = bfs(truck.pos, truck.dir_control,
                                      trash.pos, houses)
                        print(actions)
                        if not actions:
                            print('Path couldn\'t be found')
                            break

                        print('##################################################')
                        while actions:
                            action = actions.pop()
                            pygame.event.post(pygame.event.Event(
                                pygame.KEYDOWN, {'key': action}))

                            game_keys(truck, multi_trash, houses, True)
                            update_images(gameDisplay, truck, multi_trash, houses,pools,grannies)
                    else:
                        truck.space=0
                        truck.mass=0

                if (event.key == pygame.K_b):
                    trash = multi_trash[0]
                    actions = bfs(truck.pos, truck.dir_control,
                                  trash.pos, houses)
                    print(actions)
                    if not actions:
                        print('Path couldn\'t be found')
                        break

                    print('##################################################')
                    while actions:
                        action = actions.pop()
                        pygame.event.post(pygame.event.Event(
                            pygame.KEYDOWN, {'key': action}))

                        game_keys(truck, multi_trash, houses, True)
                        update_images(gameDisplay, truck, multi_trash, houses,pools,grannies)
                if (event.key == pygame.K_t):
                    solution = tsp(multi_trash, truck)
                    actions = tspmove(solution, truck, multi_trash)
                    if not actions:
                        print('Path couldn\'t be found')
                        break


                    for i in actions:
                        arr=i
                        #print(arr)
                        print('##################################################')
                        while arr:
                            action = arr.pop()
                            pygame.event.post(pygame.event.Event(
                                pygame.KEYDOWN, {'key': action}))

                            game_keys(truck, multi_trash, houses, True)
                            update_images(gameDisplay, truck, multi_trash, houses,pools,grannies)
                    truck.direction = [1, 0]
                    truck.dir_control = 0
                    truck.image = pygame.image.load('./img/truck.png')
                if not multi_trash:
                    for _ in range(7):
                        trash = Trash(game_w, game_h, grid_size)
                        trash.new_pos(truck.pos, houses, multi_trash)
                        multi_trash.append(trash)
                else:
                    pygame.event.post(event)
                    game_keys(truck, multi_trash, houses)
        update_images(gameDisplay, truck, multi_trash, houses,pools,grannies)
        clock.tick(60)


if __name__ == '__main__':
    game_loop()
pygame.quit()