game.py

"""
Snake Game Environment for AI Training

This module provides the game environment where the AI agent learns to play Snake.
It handles game logic, rendering, collision detection, and reward calculation.
"""

import pygame
import random
import numpy as np
from enum import Enum
from collections import namedtuple
from typing import Tuple, List, Optional

# Initialize pygame once
pygame.init()

# Try to load custom font, fall back to system font if not available
try:
    font = pygame.font.Font('arial.ttf', 25)
except FileNotFoundError:
    font = pygame.font.SysFont('arial', 25)


class Direction(Enum):
    """Possible movement directions for the snake"""
    RIGHT = 1
    LEFT = 2
    UP = 3
    DOWN = 4


# Lightweight point structure for coordinates
Point = namedtuple('Point', 'x, y')

# Game colors (RGB values)
class Colors:
    WHITE = (255, 255, 255)
    RED = (200, 0, 0)
    BLUE_DARK = (0, 0, 255)
    BLUE_LIGHT = (0, 100, 255)
    BLACK = (0, 0, 0)

# Game configuration
BLOCK_SIZE = 20
GAME_SPEED = 40  # Frames per second

class SnakeGameAI:
    """
    Snake Game Environment for AI Training
    
    This class manages the game state, handles AI actions, and provides rewards
    for the learning algorithm. The AI controls the snake using action arrays.
    """

    def __init__(self, width: int = 640, height: int = 480):
        """Initialize the game environment"""
        self.width = width
        self.height = height
        
        # Create the game window
        self.display = pygame.display.set_mode((self.width, self.height))
        pygame.display.set_caption('Snake AI Training')
        self.clock = pygame.time.Clock()
        
        # Initialize game state
        self.reset()

    def reset(self) -> None:
        """Reset the game to starting conditions"""
        # Snake starts moving right from the center
        self.direction = Direction.RIGHT
        
        # Create initial snake (head + 2 body segments)
        center_x, center_y = self.width // 2, self.height // 2
        self.head = Point(center_x, center_y)
        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.frame_count = 0  # Prevents infinite loops

    def _place_food(self) -> None:
        """Place food at a random location not occupied by the snake"""
        while True:
            # Generate random grid-aligned coordinates
            x = random.randint(0, (self.width - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            y = random.randint(0, (self.height - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            self.food = Point(x, y)
            
            # Make sure food doesn't appear on the snake
            if self.food not in self.snake:
                break


    def play_step(self, action: List[int]) -> Tuple[float, bool, int]:
        """
        Execute one game step with the given action
        
        Args:
            action: [straight, right, left] where exactly one element is 1
            
        Returns:
            (reward, game_over, score)
        """
        self.frame_count += 1
        
        # Handle pygame events (required to keep window responsive)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
        
        # Move the snake based on the action
        self._move_snake(action)
        self.snake.insert(0, self.head)  # Add new head
        
        # Check for game over conditions
        reward = 0
        game_over = False
        
        # Game ends if snake hits wall/itself or takes too long
        if self.is_collision() or self.frame_count > 100 * len(self.snake):
            game_over = True
            reward = -10  # Negative reward for dying
            return reward, game_over, self.score

        # Check if snake ate food
        if self.head == self.food:
            self.score += 1
            reward = 10  # Positive reward for eating
            self._place_food()  # Place new food
        else:
            self.snake.pop()  # Remove tail (snake doesn't grow)
        
        # Update display and control game speed
        self._update_display()
        self.clock.tick(GAME_SPEED)
        
        return reward, game_over, self.score

    def is_collision(self, point: Optional[Point] = None) -> bool:
        """Check if a point collides with walls or snake body"""
        if point is None:
            point = self.head
        
        # Check boundary collision
        if (point.x > self.width - BLOCK_SIZE or point.x < 0 or 
            point.y > self.height - BLOCK_SIZE or point.y < 0):
            return True
        
        # Check self-collision (don't include head in check)
        if point in self.snake[1:]:
            return True
        
        return False


    def _update_display(self) -> None:
        """Draw the current game state to the screen"""
        self.display.fill(Colors.BLACK)

        # Draw snake segments
        for segment in self.snake:
            pygame.draw.rect(self.display, Colors.BLUE_DARK, 
                           pygame.Rect(segment.x, segment.y, BLOCK_SIZE, BLOCK_SIZE))
            pygame.draw.rect(self.display, Colors.BLUE_LIGHT, 
                           pygame.Rect(segment.x + 4, segment.y + 4, 12, 12))

        # Draw food
        pygame.draw.rect(self.display, Colors.RED, 
                        pygame.Rect(self.food.x, self.food.y, BLOCK_SIZE, BLOCK_SIZE))

        # Draw score
        score_text = font.render(f"Score: {self.score}", True, Colors.WHITE)
        self.display.blit(score_text, [0, 0])
        pygame.display.flip()

    def _move_snake(self, action: List[int]) -> None:
        """
        Update snake direction and head position based on action
        
        Action encoding: [straight, turn_right, turn_left]
        """
        # Map directions in clockwise order for easy turning
        clockwise_directions = [Direction.RIGHT, Direction.DOWN, Direction.LEFT, Direction.UP]
        current_idx = clockwise_directions.index(self.direction)

        if np.array_equal(action, [1, 0, 0]):  # Go straight
            new_direction = clockwise_directions[current_idx]
        elif np.array_equal(action, [0, 1, 0]):  # Turn right
            next_idx = (current_idx + 1) % 4
            new_direction = clockwise_directions[next_idx]
        else:  # Turn left [0, 0, 1]
            next_idx = (current_idx - 1) % 4
            new_direction = clockwise_directions[next_idx]

        self.direction = new_direction

        # Update head position based on new direction
        x, y = self.head.x, self.head.y
        if self.direction == Direction.RIGHT:
            x += BLOCK_SIZE
        elif self.direction == Direction.LEFT:
            x -= BLOCK_SIZE
        elif self.direction == Direction.DOWN:
            y += BLOCK_SIZE
        elif self.direction == Direction.UP:
            y -= BLOCK_SIZE

        self.head = Point(x, y)