import numpy as np
import sys
import pygame
from pygame.locals import *

COLS = 10
ROWS = 10
array = np.zeros((ROWS, COLS))

array[0][0] = 2.0


FPS = 60
fpsClock = pygame.time.Clock()

WIN_WIDTH = 500
WIN_HEIGHT = 500
Window = pygame.display.set_mode((WIN_WIDTH, WIN_HEIGHT))
BACKGROUND = (0, 0, 0)


def draw_cell(display, position, color):
    position = (position[0] * 50 + 5, position[1] * 50 + 5)
    rect = pygame.Rect(position, (40, 40))
    pygame.draw.rect(display, color, rect)


def main():
    looping = True
    while looping:
        for event in pygame.event.get():
            if event.type == QUIT:
                print("Received QUIT event. Exiting.")
                looping = False
                pygame.quit()
                sys.exit()

        Window.fill(BACKGROUND)  # Clear display

        # Draw grid
        for row in range(0, ROWS):
            for column in range(0, COLS):
                value = array[row][column]
                if value == 0.0:
                    draw_cell(Window, (row, column), (255, 255, 255))
                elif value == 1.0:
                    draw_cell(Window, (row, column), (0, 255, 0))
                elif value == 2.0:
                    draw_cell(Window, (row, column), (255, 0, 0))

        pygame.display.update()  # Swap render buffer
        fpsClock.tick(FPS)


if __name__ == "__main__":
    pygame.init()
    main()