贪吃蛇自动游戏（GPU加速版）

import pygame
import random
import torch
import torch.nn.functional as F
from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *

# 初始化 pygame
pygame.init()

# 设置窗口大小
WIDTH, HEIGHT = 600, 400
BLOCK_SIZE = 20  # 贪吃蛇和食物的大小

# 颜色定义
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 0)

# 初始化窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.OPENGL | pygame.DOUBLEBUF)
pygame.display.set_caption("贪吃蛇游戏（GPU 加速）")

# 方向控制
UP = (0, -BLOCK_SIZE)
DOWN = (0, BLOCK_SIZE)
LEFT = (-BLOCK_SIZE, 0)
RIGHT = (BLOCK_SIZE, 0)
DIRECTIONS = [UP, DOWN, LEFT, RIGHT]

clock = pygame.time.Clock()
font = pygame.font.SysFont("Arial", 36)  # 使用 Arial 字体


def draw_text(text, color, position):
    text_surface = font.render(text, True, color)
    # Convert the text surface to a texture and display it
    text_texture = pygame.image.tostring(text_surface, "RGBA")
    glRasterPos2f(position[0], position[1])  # 设置文本位置
    glDrawPixels(text_surface.get_width(), text_surface.get_height(), GL_RGBA, GL_UNSIGNED_BYTE, text_texture)

def get_next_direction(snake, food):
    """使用 GPU 计算下一步方向，避免撞墙和自撞"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    head_x, head_y = snake[0]
    food_x, food_y = food

    directions = torch.tensor(DIRECTIONS, dtype=torch.float32, device=device)
    head_pos = torch.tensor([head_x, head_y], dtype=torch.float32, device=device)
    food_pos = torch.tensor([food_x, food_y], dtype=torch.float32, device=device)

    # 计算方向向量
    distances = torch.norm(head_pos + directions - food_pos, dim=1)
    sorted_indices = torch.argsort(distances)

    # 避免撞墙或撞自己
    for idx in sorted_indices:
        new_pos = (head_x + directions[idx][0].item(), head_y + directions[idx][1].item())
        if 0 <= new_pos[0] < WIDTH and 0 <= new_pos[1] < HEIGHT and new_pos not in snake:
            return directions[idx].tolist()

    # 如果所有方向都不可行，则尝试随机方向（不影响全局变量）
    shuffled_directions = DIRECTIONS[:]
    random.shuffle(shuffled_directions)
    for d in shuffled_directions:
        new_pos = (head_x + d[0], head_y + d[1])
        if 0 <= new_pos[0] < WIDTH and 0 <= new_pos[1] < HEIGHT and new_pos not in snake:
            return list(d)

    # 如果仍然找不到合法方向，则默认返回当前方向（防止崩溃）
    return [0, 0]


def draw_rect(color, position):
    """使用 OpenGL 绘制矩形"""
    glColor3f(*[c / 255 for c in color])
    glBegin(GL_QUADS)
    x, y = position
    glVertex2f(x, y)
    glVertex2f(x + BLOCK_SIZE, y)
    glVertex2f(x + BLOCK_SIZE, y + BLOCK_SIZE)
    glVertex2f(x, y + BLOCK_SIZE)
    glEnd()


def reset_game():
    """初始化游戏状态"""
    return [(100, 100), (80, 100), (60, 100)], \
           (random.randint(0, (WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE,
            random.randint(0, (HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE), \
           0, 10, False  # 初始蛇、食物、得分、速度、是否胜利


def main():
    while True:
        snake, food, score, speed, won = reset_game()
        running = True

        while running:
            glClear(GL_COLOR_BUFFER_BIT)
            glLoadIdentity()
            glOrtho(0, WIDTH, HEIGHT, 0, -1, 1)

            # 处理事件，防止窗口无响应
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    return

            # 计算 AI 控制的方向
            snake_direction = get_next_direction(snake, food)

            # 移动蛇
            new_head = (snake[0][0] + int(snake_direction[0]), snake[0][1] + int(snake_direction[1]))

            # 检测碰撞（边界或自身）
            if new_head[0] < 0 or new_head[0] >= WIDTH or new_head[1] < 0 or new_head[1] >= HEIGHT or new_head in snake:
                running = False
                break  # 退出循环

            # 吃食物
            if new_head == food:
                score += 1
                food = (random.randint(0, (WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE,
                        random.randint(0, (HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE)

                # 胜利条件
                if score >= 15:
                    won = True
                    running = False  # 停止游戏循环
                    break
            else:
                snake.pop()  # 没有吃到食物，移除尾部

            # 更新蛇的身体
            snake.insert(0, new_head)

            # 画蛇
            for segment in snake:
                draw_rect(GREEN, segment)

            # 画食物
            draw_rect(RED, food)

            # 更新屏幕
            pygame.display.flip()

            # 控制游戏速度
            clock.tick(speed)

        # 确保胜利或失败画面能显示
        if won:
            draw_text("You Win!", WHITE, (WIDTH // 2 - 80, HEIGHT // 2 - 20))
        else:
            draw_text("Game Over!", RED, (WIDTH // 2 - 80, HEIGHT // 2 - 20))
        draw_text(f"Your Score: {score}", WHITE, (WIDTH // 2 - 80, HEIGHT // 2 + 20))
        draw_text("Press R to Restart", WHITE, (WIDTH // 2 - 80, HEIGHT // 2 + 50))
        
        pygame.display.flip()  # 刷新画面
        
        pygame.time.wait(2000)  # 等待 2 秒，防止立刻跳转
        
        # 等待玩家按 R 重新开始
        waiting = True
        while waiting:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    return
                elif event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                    waiting = False
            pygame.time.wait(100)  # 防止 CPU 过载


# 运行游戏
main()
pygame.quit()