snake_game.py

# playable snake game
import pygame
import random
from enum import Enum
from collections import namedtuple

pygame.init()
font = pygame.font.SysFont('arial',24)

# reset
# reward
# play(action) -> direction
# game_iteration
# is_collision

class Direction(Enum):
    RIGHT = 1
    LEFT = 2
    UP = 3
    DOWN = 4

Point = namedtuple('Point',['x','y'])

#rgb color
WHITE = (255,255,255)
RED = (200,0,0)
BLUE1 = (0,0,255)
BLUE2 = (0,100,255)
BLACK = (0,0,0)

BLOCK_SIZE = 20
SPEED = 10

class SnakeGame:
    def __init__(self, w = 640, h = 480):
        self.w = w
        self.h = h

        #init display
        self.display = pygame.display.set_mode((self.w,self.h))
        pygame.display.set_caption('Snake')
        self.clock = pygame.time.Clock()
        self._reset()

        #init game state
    def _reset(self):
        self.direction = Direction.RIGHT
        self.head = Point(self.w/2, self.h/2)
        self.snake = [self.head,Point(self.head.x - BLOCK_SIZE, self.head.y),Point(self.head.x - 2*BLOCK_SIZE, self.head.y)]

        self.score = 0
        self.food = None
        self._place_food()
        self.game_iteration = 0

    def play_step(self):
        #1.collect user input
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_a:
                   self.direction = Direction.LEFT
                elif event.key == pygame.K_d:
                   self.direction = Direction.RIGHT
                elif event.key == pygame.K_w:
                   self.direction = Direction.UP 
                elif event.key == pygame.K_s:
                   self.direction = Direction.DOWN  
        #2.move
        self._move(self.direction)
        self.snake.insert(0,self.head)

        #3.check if game over
        game_over = False
        if self._is_collision():
            game_over = True
            return game_over, self.score

        #4. place new food or just move
        if self.head == self.food:
            self.score += 1
            self._place_food()
        else:
            self.snake.pop()

        #5.update ui and clock
        self._update_ui()
        self.clock.tick(SPEED)
        #6.return game over and score

        return game_over,self.score
    
    def _place_food(self):
        x = random.randint(0, (self.w-BLOCK_SIZE)//BLOCK_SIZE) *BLOCK_SIZE
        y = random.randint(0, (self.h-BLOCK_SIZE)//BLOCK_SIZE) *BLOCK_SIZE
        self.food = Point(x,y)
        if self.food in self.snake:
            self._place_food()

    def _is_collision(self):
        #hit boundary
        if self.head.x > self.w - BLOCK_SIZE or self.head.x < 0 or self.head.y < 0 or self.head.y > self.h - BLOCK_SIZE:
            return True
        #hit it self
        if self.head in self.snake[1:]:
            return True
        return False

    def _move(self,direction):
        x = self.head.x
        y = self.head.y
        if direction == Direction.RIGHT:
            x += BLOCK_SIZE
        if direction == Direction.LEFT:
            x -= BLOCK_SIZE
        if direction == Direction.UP:
            y -= BLOCK_SIZE
        if direction == Direction.DOWN:
            y += BLOCK_SIZE
        
        self.head = Point(x,y)

    def _update_ui(self):
        self.display.fill(BLACK)
        for point in self.snake:
            pygame.draw.rect(self.display,BLUE1, pygame.Rect(point.x, point.y, BLOCK_SIZE, BLOCK_SIZE))
            pygame.draw.rect(self.display,BLUE2, pygame.Rect(point.x+4, point.y+4, BLOCK_SIZE-8, BLOCK_SIZE-8))
        
        pygame.draw.rect(self.display,RED, pygame.Rect(self.food.x,self.food.y,BLOCK_SIZE,BLOCK_SIZE))
        text = font.render("Score: " + str(self.score),True,WHITE)
        self.display.blit(text,[0,0])
        pygame.display.flip()


if __name__ == '__main__':
    game = SnakeGame()

    #gameloop
    while True:
        game_over, score = game.play_step()

        #break if game over
        if game_over:
            break

    print('Final score',score)

    pygame.quit()