2019SZI-Projekt/UI/grid.py

import pygame as pg
import numpy as np
import random as rd
from os import listdir
from os.path import isfile, join
# from Logic.TrashRecognition.ImageClassification import classify

# MODULE LEVEL VARIABLES
recognized_trash = {

}
# trash_files = classify()
########################


class Grid:

    def __init__(self, cols: int, rows: int):
        self.table = [[Node(row, col)
                       for col in range(cols)]
                      for row in range(rows)]
        self.cols = cols
        self.rows = rows

    def draw_map(self, screen: "PyGame screen"):
        """Draws whole map"""
        screen.fill(Node.BLACK)
        for row in self.table:
            for node in row:
                node.draw(screen)

    def change_field(self, row: int, col: int, f_type: int):
        self.table[row][col].change_field_type(f_type)
    
    def generate_trash(self,row: int, col: int):
        self.table[row][col].generate_trash()

    def draw_node(self, screen, row: int, col: int):
        self.table[row][col].draw(screen)
    
    def get_trash_possition(self, trash: int):
        trash_possition = []
        for row in self.table:
            for node in row:
                if node.field_type == trash and node.house.empty == False:
                    trash_possition.append((node.row,node.col))
        return trash_possition
    
    def garbage_to_collect(self):
        garbage = []
        field_types = []
        d = rd.randint(1,7)
        day = self.table[0][0].house.get_day_of_week(d)
        for index, x in enumerate(day):
            if index != 0 and index != 1:
                garbage.append(x)
        print("\n Today is:", day[1], ", garbage to collect: ",end = '')
        for x in garbage:
            print(x[2],", ",end='')
            field_types.append(x[0])
        return field_types




class House:
    # define some basic colors
    BLACK = (0, 0, 0)
    WHITE = (255, 255, 255)
    GREEN = (0, 255, 0)
    SKYBLUE = (0,191,255)
    RED = (255, 0, 0)
    BLUE = (0, 0, 255)
    ORANGE = (255, 165, 0)
    PINK = (199,21,133)
    GREY = (192,192,192)

    #define trash 
    trash_dict = {
        "paper": (5, WHITE, "paper"),
        "glass": (6, SKYBLUE, "glass"),
        "metal": (7, GREY, "metal"),
        "plastic": (8, ORANGE, "plastic")
    }
    #define days of the week
    MONDAY = (1, "Monday", trash_dict["paper"], trash_dict["metal"])
    TUESDAY = (2, "Tuesday", trash_dict["glass"])
    WEDNESDAY=(3, "Wednesday", trash_dict["plastic"], trash_dict["metal"])
    THURSDAY = (4, "Thursday", trash_dict["glass"])
    FRIDAY = (5, "Friday", trash_dict["paper"], trash_dict["metal"])
    SATURDAY = (6, "Saturday", trash_dict["plastic"])
    SUNDAY = (7, "Sunday", trash_dict["metal"])
    DAYS = [MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY, SATURDAY, SUNDAY]


    def __init__(self):
        self.empty = True
        self.trash = None
        self.trash_file = None

            
    def find_trash_file(self):
        from os.path import sep # culture and os invariant separator
        
        file_names = [(join(f"Images{sep}TestImages", f))
             for f in listdir(f"Images{sep}TestImages") 
             if isfile(join(f"Images{sep}TestImages", f))]

        file_name = file_names[rd.randint(0,len(file_names)) - 1]

        from Logic.TrashRecognition.ImageClassification import classify_file

        if file_name in recognized_trash:
            rt = recognized_trash[file_name]
            return (file_name, rt[0], rt[1])
        else:
            classification = classify_file(file_dir=file_name)
            recognized_trash[file_name] = (classification[1], classification[2])
            return classification 

    def generate_trash(self):
        self.empty = False
        
        classification = self.find_trash_file()
        self.trash = self.trash_dict[classification[1]]
        self.trash_file = classification[0]

    def get_day_of_week(self, d: int):
        for day in self.DAYS:
            if day[0] == d:
                return day



class Node:
    # define rectangles dimensions
    r_width = 20
    r_height = 20
    r_margin = 5

    # define some basic colors
    # TODO: change to Enum
    BLACK = (0, 0, 0)
    WHITE = (255, 255, 255)
    GREEN = (0, 255, 0)
    SKYBLUE = (0,191,255)
    RED = (255, 0, 0)
    BLUE = (0, 0, 255)
    ORANGE = (255, 165, 0)
    PINK = (199,21,133)
    GREY = (192,192,192)


    def __init__(self, row: int, col: int,
                 field_type: int = 0):
        self.row = row
        self.col = col
        self.field_type = field_type
        self.house = House()



    def draw(self, screen):
        color = self.get_field_color()
        col = self.col
        row = self.row
        width = self.r_width
        height = self.r_height
        margin = self.r_margin
        # rect -> (left, top, width, height)
        # draw.rect(surface, color, rect, margin)
        pg.draw.rect(screen, color,
                     ((col * (width + margin)) + margin,
                      (row * (height + margin)) + margin,
                      width, height))
        pg.display.flip()  # refresh screen
    
    def generate_trash(self):
        self.house.generate_trash()
        self.field_type = self.house.trash[0]


    def change_field_type(self, field_type: int):
        self.field_type = field_type

    def get_field_color(self) -> tuple:
        """Gets the color tuple of field"""
        #base color
        if self.field_type == 0:
            return self.GREEN
        #truck color
        elif self.field_type == 1:
            return self.RED
        #end point
        elif self.field_type == 2:
            return self.BLUE
        #obstacles color
        elif self.field_type == 3:
            return self.ORANGE
        #path color
        elif self.field_type == 4:
            return self.PINK
        #paper
        elif self.field_type == 5:
            return self.WHITE
        #glass
        elif self.field_type == 6:
            return self.SKYBLUE
        #metal 
        elif self.field_type == 7:
            return self.GREY
        #plastic 
        elif self.field_type == 8:
            return self.ORANGE