import pygame
import numpy as np

def main_loop():
    BLACK = (0, 0, 0)
    WHITE = (200, 200, 200)
    GREEN = (141, 182, 0)
    RED = (255, 0, 0)
    WIDTH = 20
    HEIGHT = 20
    SIZE = 11

    MARGIN = 5

    grid = []
    snake_tails = []

    for row in range(SIZE+1):

        grid.append([])
        for column in range(SIZE+1):
            grid[row].append(0)


    coord = (5, 6)

    for i in range(3):
        snake_tails.append([coord[0], coord[1] + i])
        grid[coord[0]][coord[1] + i] = 1
    # print(snake_tails)
    pygame.init()


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


    pygame.display.set_caption("Array Backed Grid")


    done = False

    clock = pygame.time.Clock()


    h=0
    v=0
    print(snake_tails)
    while not done:
        for i in pygame.event.get():
            if i.type == pygame.QUIT:
                done = True

            elif i.type == pygame.KEYDOWN:
                if i.key == pygame.K_LEFT:
                    h=-1
                    v=0
                elif i.key == pygame.K_RIGHT:
                    h=1
                    v=0
                elif i.key == pygame.K_DOWN:
                    v=1
                    h=0
                elif i.key == pygame.K_UP:
                    v=-1
                    h=0


        if h + v != 0:
            # remove tail segment
            tail = snake_tails[0]
            grid[tail[0]][tail[1]] = 0
            snake_tails[:] = snake_tails[1:]
            # append segment to head
            head = snake_tails[-1]
            new_segment = [head[0]+v, head[1]+h]
            print(snake_tails, 'head', head, 'new', new_segment)
            new_segment_ok = 0 <= new_segment[0] < SIZE and 0 <= new_segment[1] < SIZE
            if (not new_segment_ok):
                done = True
                break
            snake_tails.append(new_segment)
            head = new_segment
            grid[head[0]][head[1]] = 1
            print(snake_tails)


        for row in range(SIZE+1):
            for column in range(SIZE+1):
                color = WHITE
                if grid[row][column] == 1:
                    color = GREEN

                pygame.draw.rect(screen,
                                color,
                                [(MARGIN + WIDTH) * column + MARGIN,
                                (MARGIN + HEIGHT) * row + MARGIN,
                                WIDTH,
                                HEIGHT])


        clock.tick(10)

        pygame.display.update()
        # screen.fill(WHITE)

    pygame.quit()

main_loop()