import pygame
import random
import numpy as np
import time
import collections
import alg
from models.dumpster import trash
from models.Garbagetruck import GarbageTruck


#heuristics for finding the closest dumpster
def closest(GT, D):
    point = np.column_stack((np.abs(GT[0]-D[:,0])+np.abs(GT[1]-D[:,1]),D))
    sorted_list = sorted(point, key=lambda x:x[0])
    return sorted_list[0][1:]


def display(grid, GT, a):
    #find closest dumpster and the path
    ZZ = closest(list(GT.coor), a)
    b = alg.aStar(grid, tuple(GT.coor), tuple(ZZ))

    # go through the route to the closest dumpster
    for i in b:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                done = True
        GT.turn(i)
        GT.move()

        # "Speed" of the Garbage Truck
        clock.tick(4)

        grid=drawGrid(grid, 20)
        screen.blit(GT.picture, ((WIDTH) * GT.coor[0], (HEIGHT) *GT.coor[1]))
        pygame.display.flip()
    return ZZ



#creating dumpsters
def createDumpstersAndDump(size,number):
    points=[]
    for x in range (0,size):
        for y in range (0,size):
            points.append([x,y])

    wol=random.sample(points,number)
    return (wol)


def drawGrid(grid, size):
    for row in range(size):
        for column in range(size):
            #colouring dumpsters
            img=grassImg
            if grid[row][column] == 1:
                #plastic
                img = garbageY
            elif grid[row][column] == 2:
                #paper
                img = garbageG
            elif grid[row][column] == 3:
                #glass
                img = garbageR
            elif grid[row][column] == 4:
                #organic
                img = garbageBL
            elif grid[row][column] == 5:
                #mixede
                img = garbageBLK
            elif row == TheDump[:,0] and column == TheDump[:,1]:
                grid[row][column] = 0
                img = recycImg
            screen.blit(img, ((WIDTH) * row, (HEIGHT) * column))

    return grid

# 700:20=35
WIDTH = 35
HEIGHT = 35

# Creation of a two dimensional grid
grid = []
for row in range(20):
    grid.append([])
    for column in range(20):
        grid[row].append(10)


# number of dumpsters
size=20

#randomizing coordinates, initializing the Garbage Dump
a = createDumpstersAndDump(20,size+1)
TheDump = np.reshape(a[-1], (1, 2))
a = np.reshape(a[:-1], (size, 2))

#initiating the coordinates for the garbage truck
GT = GarbageTruck()
coor = GT.start()


#Initializing dumpsters
dumpsters = []
for i in range(0,size):
    s = trash()
    s.giveID(i)
    s.throwAway()
    s.colour()
    s.xy = a[i]
    dumpsters.append(s)

#colouring trashcans
buff = 0
for x in a:
    grid[x[0]][x[1]] = dumpsters[buff].color
    buff += 1

# Initializing, sizing the window & title
pygame.init()
winsize=700
WINDOW_SIZE = [winsize, winsize]
screen = pygame.display.set_mode(WINDOW_SIZE)
pygame.display.set_caption("Intelligent Garbage Truck")

# Loop until user clicks close
done = False

# Managing refreshing speed
clock = pygame.time.Clock()

#resources
grassImg = pygame.image.load('./images/grass.png')
garbageY = pygame.image.load('./images/garbageyellow.png')
garbageR = pygame.image.load('./images/garbagered.png')
garbageBL = pygame.image.load('./images/garbageblue.png')
garbageBLK = pygame.image.load('./images/garbageblack.png')
garbageG = pygame.image.load('./images/garbagegreen.png')
recycImg = pygame.image.load('./images/recycling.png')

all = a
while not done:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            done = True

    # Draw the grid
    grid=drawGrid(grid, 20)

    #move and display
    ZZ = display(grid, GT, a)


    for x in all:
        if(x[0] == GT.coor[0] and x[1] == GT.coor[1]):

            #recognize which dumpster are we visiting
            Xz=np.where(all==x)[0]
            x =[item for item, count in collections.Counter(Xz).items() if count > 1]

            if GT.collectingTrash(dumpsters, x[0]):
                dumpsters[x[0]].empty()
            elif GT.amIFull():
                print("DUMPSTER TRUCK FULL")
            else:
                print("No more space for this kind of trash")

    # remove visited dumpster
    a = a[~((a[:,0] ==ZZ[0]) & (a[:,1] ==ZZ[1]))]
    print("Dumpsters left: ",len(a))

    # visit Garbage Dump at the end
    if len(a) == 0:
        print("Going to the Garbage Dump")
        display(grid, GT, TheDump)
        done = True

pygame.quit()