import pygame
import random
import numpy as np
import time
import collections
import alg
from models.dumpster import trash
from models.Garbagetruck import GarbageTruck
import numbering
from sklearn.datasets import load_digits
import garbageDumpSorting as gds
import time
from models.garbageDump import Dump

#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:
                #cardboard
                img = garbageBL
            elif grid[row][column] == 5:
                #metal
                img = garbageBLK
            elif row == dump.coor[:,0] and column == dump.coor[:,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)
dump = Dump()
dump.coor = np.reshape(a[-1], (1, 2))
a = np.reshape(a[:-1], (size, 2))


#Klaudia Przybylska subproject part
gds.createSets()
# gds.processTrainData()
# gds.processTestData()
clf = gds.trainAndTest()
time.sleep(5)



#Patryk Krawiec subproject part
digitsy=np.load("Ytest.npy")
digitsx=np.load("Xtest.npy")
images_and_labels=list(zip(digitsx,digitsy))
num,xnum,ynum=numbering.num(20,images_and_labels)

paying_customers=num[3:]
paying_customers = list(map(int, paying_customers))

#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]
    s.numberx=xnum[i]
    s.numbery=ynum[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]
            predicted=numbering.pred(dumpsters[x[0]].numberx,dumpsters[x[0]].numbery)

            #recognizing whether a house is our client and trying to collect trash
            if(predicted not in paying_customers):
                print("House number is: ",predicted," but it is not our customer")
            else:
                print("House number is: ",predicted)
                GT.collectingTrash(dumpsters, x[0])



    # 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, dump.coor)
        done = True
        print("Checking garbage inside the truck: ")
        gds.sortDump(GT.plastic, GT.paper, GT.metal, GT.cardboard, GT.glass, clf, GT, dump)
        print("Garbage sorted: ")
        print("Paper: " + str(dump.paper))
        print("Cardboard: " + str(dump.cardboard))
        print("Metal: " + str(dump.metal))
        print("Plastic: " + str(dump.plastic))
        print("Glass: " + str(dump.glass))



pygame.quit()