import pygame
import numpy as np
import math

     ########################
     ###        WS        ###
     ########################
# For CNN:

import keras
from keras.preprocessing import image
from keras.models import Sequential
from keras.layers import Convolution2D
from keras.layers import MaxPooling2D
from keras.layers import Flatten
from keras.layers import Dense
 

#initializing:
classifier = Sequential()
#Convolution:
classifier.add(Convolution2D(32, (3, 3), input_shape =(256, 256, 3), activation = "relu"))
#Pooling:
classifier.add(MaxPooling2D(pool_size = (2,2)))

# Adding a second convolutional layer
classifier.add(Convolution2D(32, 3, 3, activation = 'relu'))
classifier.add(MaxPooling2D(pool_size = (2, 2)))

#Flattening:
classifier.add(Flatten())

#Fully connected layers::
classifier.add(Dense(units = 128, activation = "relu"))
classifier.add(Dense(units = 3, activation = "softmax"))

# loading weigjts:
classifier.load_weights('s444523/best_model_weights2.h5')
#Making CNN:
classifier.compile(optimizer = "adam", loss = "categorical_crossentropy", metrics = ["accuracy"])


     ########################
     ###        WS        ###
     ########################
# Colors:
# Define some colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 240)

#Width and Height of each square:
WIDTH = 20
HEIGHT = 20

#Margin:
MARGIN = 5
grid = [[0 for x in range(16)] for y in range(16)]

def change_value(i, j, width, n):
    for r in range (i, i+width):
        for c in range (j, j+width):
            grid[r][c] = n
    
# the class "Table"        
class Table:
    def __init__(self, coordinate_i, coordinate_j, state = 0):
        self.coordinate_i = coordinate_i
        self.coordinate_j = coordinate_j
        self.state = state
        change_value(coordinate_i, coordinate_j, 2, 1)
    def get_destination_coor(self):
        return [self.coordinate_i, self.coordinate_j-1]
    
     ########################
     ###        WS        ###
     ########################
     
    # The finction "state of meal" chooses a photo of a plate at the given table.
    def state_of_meal(self):
        ## !!!!!!###
        num = np.random.randint(67, 100)
                ## !!!!!!###

        if num<=67:
            img_name = 'plates/{}.png'.format(num)
        else:
            img_name = 'plates/{}.jpg'.format(num)
        return img_name
    
    # The function "change state" changes the value of the state variable.
    # It informs, whether the client are waiting for the waiter to make an order
    # (0 - empty plates) are eating (2 - full plates) or are waiting for the 
    # waiter to get a recipt (1 - dirty plates)
    
    def change_state(self, st):
        self.state = st
        
    ########################
    ###       /WS        ###
    ########################
    

class Kitchen:
    def __init__(self, coordinate_i, coordinate_j):
        self.coordinate_i = coordinate_i
        self.coordinate_j = coordinate_j
        change_value(coordinate_i, coordinate_j, 3, 2)
        
class Agent:
    def __init__(self,orig_coordinate_i, orig_coordinate_j):
        self.orig_coordinate_i = orig_coordinate_i
        self.orig_coordinate_j = orig_coordinate_j
        self.state = np.array([1,2])
        change_value(orig_coordinate_j, orig_coordinate_j, 1, 3)
        self.state_update(orig_coordinate_i, orig_coordinate_j)
        self.previous_grid = np.array([-1, -1])
        
    def state_update(self, c1, c2):
        self.state[0] = c1
        self.state[1] = c2
        
    def leave(self):
        change_value(self.state[0], self.state[1], 1, 0)
    
    def previous_grid_update(self):
        self.previous_grid[0] = self.state[0]
        self.previous_grid[1] = self.state[1]
    
    def move_to(self, nextPos):
        self.previous_grid_update()
        self.leave()
        self.state_update(x + nextPos[0], y + nextPos[1])
        change_value(self.state[0], self.state[1], 1, 3)
   
       ########################
       ###        WS        ###
       ########################
       
       #The function "define_table" serches for the apropriate table in the 
       # table_dict (to enable the usage of class attributes and methods)
    def define_table(self, t_num):
        t_num = 'table{}'.format(t_num)
        t_num = table_dict[t_num]
        return t_num

       # The function "check_plates" uses the pretrained CNN model to classify 
       # the plate on the table as empty(0), full(2) or dirty(1)
    def check_plates(self, table_number):
        table = self.define_table(table_number)
        plate = table.state_of_meal()        
        plate= image.load_img(plate, target_size = (256, 256))
        plate = image.img_to_array(plate)
        plate = np.expand_dims(plate, axis = 0)
        result = classifier.predict(plate)[0]
        print (result)
        if result[1] == 1:
            result[1] = 0
            x = int(input("Excuse me, have You done eating? 1=Yes, 2 = No \n"))
            result[x] = 1
        for i, x in enumerate(result):
            if result[i] == 1:
                table.change_state(i)

       ########################
       ###       /WS        ###
       ########################
    # check the next grid is not the previous grid to prevent the loop
    def next_is_previous(self, x, y):
        return np.array_equal(self.previous_grid, np.array([x, y]))

def check_done():
    for event in pygame.event.get():        # Checking for the event
        if event.type == pygame.QUIT:       # If the program is closed:
            return True                     # To exit the loop

def draw_grid():
    for row in range(16):                   # Drawing the grid
        for column in range(16):
            color = WHITE
            if grid[row][column] == 1:
                color = GREEN
            if grid[row][column] == 2:
                color = RED
            if grid[row][column] == 3:
                color = BLUE
            pygame.draw.rect(screen,
                             color,
                             [(MARGIN + WIDTH) * column + MARGIN,
                              (MARGIN + HEIGHT) * row + MARGIN,
                              WIDTH,
                              HEIGHT]) 

# calculate the distance between two points 
def distance(point1, point2):
    return math.sqrt((point2[0] - point1[0])**2 + (point2[1] - point1[1])**2)

## default positions of the agent:
x = 12
y = 12
agent = Agent(x, y)

table1 = Table(2, 2)
table2 = Table (2,7)
table3 = Table(2, 12)
table4 = Table(7, 2)
table5 = Table(7, 7)
table6 = Table(7, 12)
table7 = Table(12, 2)
table8 = Table(12, 7)


################### WS #####################
# I added the dict to loop through tables.
table_dict = {"table1":table1, "table2":table2, "table3":table3,"table4":table4,
              "table5":table5,"table6":table6,"table7":table7,"table8":table8
              }
################### WS #####################

#class Kitchen:
kitchen = Kitchen(13, 13)

# destination array
destination_tables = []


pygame.init()
WINDOW_SIZE = [405, 405]
screen = pygame.display.set_mode(WINDOW_SIZE)

pygame.display.set_caption("Waiter_Grid3")

done = False

clock = pygame.time.Clock()
 
# -------- Main Program Loop -----------
while not done:
    screen.fill(BLACK)                      # Background color
    
    draw_grid()
    done = check_done()
    for value in table_dict.values(): destination_tables.append(value.get_destination_coor())
    # We need to define the number of the table by which we are in:
    
    num_of_table = 1
    while len(destination_tables) != 0:
        
        # set the first element(table) in array as currDestination
        currDestination = destination_tables[0]
        # from kitchen to table
        while agent.state[0] != currDestination[0] or agent.state[1] != currDestination[1]:

            #///////////////////////////////////////
            x = agent.state[0]
            y = agent.state[1]
            
            # set a huge default number
            minDis = 9999
            nextPos = []
            # check whether the agent goes left
            if y-1 >= 0 and grid[x][y-1] != 1 and not agent.next_is_previous(x, y-1):
                minDis = distance([x, y-1], currDestination)
                nextPos = [0, -1]  # means go left
            
            # check whether the agent goes right
            if y+1 <= 15 and grid[x][y+1] != 1 and grid[x][y+1] != 2 and not agent.next_is_previous(x, y+1):  
                d = distance([x, y+1], currDestination)
                if d < minDis:
                    minDis = d
                    nextPos = [0, 1]  # means go right
            
            # check whether the agent goes up
            if x-1 >= 0 and grid[x-1][y] != 1 and not agent.next_is_previous(x-1, y):
                d = distance([x-1, y], currDestination)
                if d < minDis:
                    minDis = d
                    nextPos = [-1, 0]  # means go up
                    
            # check whether the agent goes down
            if x+1 <= 15 and grid[x+1][y] != 1 and grid[x+1][y] != 2 and not agent.next_is_previous(x+1, y):
                d = distance([x+1, y], currDestination)
                if d < minDis:
                    minDis = d
                    nextPos = [1, 0]  # means go down
            
#            print(agent.previous_grid)
            agent.move_to(nextPos)
            #////////////////////////////////////////////////
            
            pygame.time.delay(100)
            screen.fill(BLACK)                      # Background color
            draw_grid()                             # Drawing the grid
            clock.tick(60)                           # Limit to 60 frames per second
            pygame.display.flip()                    # Updating the screen
        
        
        ########################
        ###        WS        ###
        ########################
        #pygame.time.delay(100)
        print("I'm at a table no. {}".format(num_of_table))
        ## Checking at what state are the plates:
        agent.check_plates(num_of_table)
        num_of_table +=1
        
        ########################
        ###       /WS        ###
        ########################
        # set the kitchen as currDestination 
        currDestination = [13, 12]
        # from table to kitchen
        while agent.state[0] != currDestination[0] or agent.state[1] != currDestination[1]:

            #///////////////////////////////////////
            x = agent.state[0]
            y = agent.state[1]
            
            # set a huge default number
            minDis = 9999
            nextPos = []
            # check whether the agent goes left
            if y-1 >= 0 and grid[x][y-1] != 1 and not agent.next_is_previous(x, y-1):
                minDis = distance([x, y-1], currDestination)
                nextPos = [0, -1]  # means go left
            
            # check whether the agent goes right
            if y+1 <= 15 and grid[x][y+1] != 1 and grid[x][y+1] != 2 and not agent.next_is_previous(x, y+1):  
                d = distance([x, y+1], currDestination)
                if d < minDis:
                    minDis = d
                    nextPos = [0, 1]  # means go right
            
            # check whether the agent goes up
            if x-1 >= 0 and grid[x-1][y] != 1 and grid[x-1][y] != 2 and not agent.next_is_previous(x-1, y):
                d = distance([x-1, y], currDestination)
                if d < minDis:
                    minDis = d
                    nextPos = [-1, 0]  # means go up
                    
            # check whether the agent goes down
            if x+1 <= 15 and grid[x+1][y] != 1 and grid[x+1][y] != 2 and not agent.next_is_previous(x+1, y):
                d = distance([x+1, y], currDestination)
                if d < minDis:
                    minDis = d
                    nextPos = [1, 0]  # means go down
            
            agent.move_to(nextPos)
            #////////////////////////////////////////////////
            
            pygame.time.delay(100)
            screen.fill(BLACK)                      # Background color
            draw_grid()                             # Drawing the grid
            clock.tick(60)                           # Limit to 60 frames per second
            pygame.display.flip()                    # Updating the screen
            
        
        destination_tables = destination_tables[1:]    # remove the first element in the list
        # After each fool loop, we can quit the program:.
        if len(destination_tables) == 0:
            play_again = 1
            play_again = int(input("Exit? 0=No, 1=Yes \n"))
            if play_again:
                pygame.quit()
   
    
pygame.quit()