import pygame
import random
import pandas as pd
from sklearn import tree
from sklearn.preprocessing import LabelEncoder
import graphviz


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

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))
DIRT_IMG = pygame.image.load("dirt.jpg")
DIRT = pygame.transform.scale(DIRT_IMG, (50, 50))
GRASS_IMG = pygame.image.load("grass.png")
GRASS = pygame.transform.scale(GRASS_IMG, (50, 50))
SAND_IMG = pygame.image.load("sand.jpeg")
SAND = pygame.transform.scale(SAND_IMG, (50, 50))
COBBLE_IMG = pygame.image.load("cobble.jpeg")
COBBLE = pygame.transform.scale(COBBLE_IMG, (50, 50))
FPS = 10
FIELDCOUNT = 16
FIELDWIDTH = 50


class Agent:
    def __init__(self, rect, direction):
        self.rect = rect
        self.direction = direction


def randomize_map():        # tworzenie mapy z losowymi polami
    field_array_1 = []
    field_array_2 = []
    field_priority = []
    for i in range(16):
        temp_priority = []
        for j in range(16):
            if i in (0, 1) and j in (0, 1):
                field_array_2.append(GRASS)
                temp_priority.append(1)
            else:
                prob = random.uniform(0, 100)
                if 0 <= prob <= 12:
                    field_array_2.append(COBBLE)
                    temp_priority.append(3)
                elif 12 < prob <= 24:
                    field_array_2.append(SAND)
                    temp_priority.append(2)
                else:
                    field_array_2.append(GRASS)
                    temp_priority.append(1)
        field_array_1.append(field_array_2)
        field_array_2 = []
        field_priority.append(temp_priority)
    return field_array_1, field_priority


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():
    train_data = pd.read_csv('./data_set.csv')
    attributes = train_data.drop('collect', axis='columns')
    e_type = LabelEncoder()
    attributes['type_num'] = e_type.fit_transform(attributes['garbage_type'])
    attr_encoded = attributes.drop(['garbage_type'], axis='columns')
    attr_names = ['fuel','distance','space_occupied','days_since_last_collection','paid_on_time','odour_intensity','garbage_weight', 'garbage_type']
    label_names = ['collect', 'no-collect']
    label = train_data['collect']
    print(attr_encoded)
    print(label)
    classifier = tree.DecisionTreeClassifier()
    classifier.fit(attr_encoded, label)
    dot_data = tree.export_graphviz(classifier, out_file=None, feature_names=attr_names, class_names=label_names)
    graph = graphviz.Source(dot_data)
    graph.render('collect')
    clock = pygame.time.Clock()
    run = True
    x, y = [0, 0]
    agent = GarbageTruck(0, 0, pygame.Rect(x, y, 50, 50), 0)   # tworzenie pola dla agenta
    fields, priority_array = randomize_map()
    final_x, final_y = [100, 300]
    while run:
        clock.tick(FPS)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
        # keys_pressed = pygame.key.get_pressed()
        draw_window(agent, fields, False)   # false = kierunek east (domyslny), true = west
        steps = astar(State(None, None, x, y, 'E', priority_array[0][0], heuristicfn(x, y, final_x, final_y)), final_x, final_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.5)
    
        while True:
            pass  

    pygame.quit()


if __name__ == "__main__":
    main()