import pygame
import sys

# 初始化pygame
pygame.init()

# 游戏常量设置
SCREEN_SIZE = 600       # 窗口大小
GRID_SIZE = 30          # 每个格子的大小
LINE_WIDTH = 2          # 棋盘线宽度
BOARD_SIZE = 19         # 棋盘大小(19x19)
PIECE_RADIUS = 14       # 棋子半径
PIECE_WIDTH = 2         # 棋子边框宽度

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BG_COLOR = (240, 220, 180)  # 棋盘背景色(木纹色)
LINE_COLOR = (100, 100, 100)# 棋盘线颜色

# 计算棋盘起始位置（居中显示）
START_X = (SCREEN_SIZE - (BOARD_SIZE - 1) * GRID_SIZE) // 2
START_Y = (SCREEN_SIZE - (BOARD_SIZE - 1) * GRID_SIZE) // 2

# 创建游戏窗口
screen = pygame.display.set_mode((SCREEN_SIZE, SCREEN_SIZE))
pygame.display.set_caption("五子棋")

# 游戏模式：0表示双人对战，1表示AI对战
game_mode = 1  # 默认AI对战
mode_selected = False  # 是否已选择游戏模式

# 初始化棋盘数据，0表示空，1表示黑子，2表示白子
board = [[0 for _ in range(BOARD_SIZE)] for _ in range(BOARD_SIZE)]
current_player = 1  # 1表示黑方，2表示白方
game_over = False
winner = None

# AI评估权重
WEIGHTS = {
    'win': 100000,        # 获胜
    'block_win': 50000,    # 阻挡对方获胜
    'four': 10000,         # 四子连珠
    'block_four': 5000,    # 阻挡对方四子
    'three': 1000,         # 三子连珠
    'block_three': 500,    # 阻挡对方三子
    'two': 100,            # 二子连珠
    'block_two': 50,       # 阻挡对方二子
    'one': 10,             # 一子
    'block_one': 5         # 阻挡对方一子
}

def draw_board():
    """绘制棋盘"""
    # 填充背景色
    screen.fill(BG_COLOR)
    
    # 绘制棋盘网格线
    for i in range(BOARD_SIZE):
        # 竖线
        pygame.draw.line(
            screen, LINE_COLOR,
            (START_X + i * GRID_SIZE, START_Y),
            (START_X + i * GRID_SIZE, START_Y + (BOARD_SIZE - 1) * GRID_SIZE),
            LINE_WIDTH
        )
        # 横线
        pygame.draw.line(
            screen, LINE_COLOR,
            (START_X, START_Y + i * GRID_SIZE),
            (START_X + (BOARD_SIZE - 1) * GRID_SIZE, START_Y + i * GRID_SIZE),
            LINE_WIDTH
        )
    
    # 绘制棋子
    for row in range(BOARD_SIZE):
        for col in range(BOARD_SIZE):
            if board[row][col] != 0:
                center_x = START_X + col * GRID_SIZE
                center_y = START_Y + row * GRID_SIZE
                color = BLACK if board[row][col] == 1 else WHITE
                # 绘制棋子
                pygame.draw.circle(screen, color, (center_x, center_y), PIECE_RADIUS)
                # 绘制棋子边框（增加视觉效果）
                pygame.draw.circle(screen, BLACK, (center_x, center_y), PIECE_RADIUS, PIECE_WIDTH)

def get_click_pos(pos):
    """将鼠标点击位置转换为棋盘坐标"""
    x, y = pos
    # 计算所在格子
    col = (x - START_X + GRID_SIZE // 2) // GRID_SIZE
    row = (y - START_Y + GRID_SIZE // 2) // GRID_SIZE
    
    # 检查是否在棋盘范围内
    if 0 <= row < BOARD_SIZE and 0 <= col < BOARD_SIZE:
        return (row, col)
    return None

def check_win(row, col):
    """检查落子后是否获胜"""
    player = board[row][col]
    directions = [
        (0, 1),   # 横向
        (1, 0),   # 纵向
        (1, 1),   # 正对角线
        (1, -1)   # 反对角线
    ]
    
    for dr, dc in directions:
        count = 1
        # 正方向检查
        r, c = row + dr, col + dc
        while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and board[r][c] == player:
            count += 1
            r += dr
            c += dc
        
        # 反方向检查
        r, c = row - dr, col - dc
        while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and board[r][c] == player:
            count += 1
            r -= dr
            c -= dc
        
        # 五子连珠则获胜
        if count >= 5:
            return True
    return False

def evaluate_position(row, col, player):
    """评估指定位置对玩家的价值"""
    if board[row][col] != 0:
        return -1000000  # 已经有棋子的位置
    
    score = 0
    opponent = 1 if player == 2 else 2
    
    # 模拟落子
    board[row][col] = player
    
    # 检查是否获胜
    if check_win(row, col):
        score += WEIGHTS['win']
    
    # 评估各个方向的连珠情况
    directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
    for dr, dc in directions:
        # 检查当前玩家的连珠
        count = 1
        # 正方向
        r, c = row + dr, col + dc
        while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and board[r][c] == player:
            count += 1
            r += dr
            c += dc
        # 反方向
        r, c = row - dr, col - dc
        while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and board[r][c] == player:
            count += 1
            r -= dr
            c -= dc
        
        # 根据连珠数量加分
        if count == 4:
            score += WEIGHTS['four']
        elif count == 3:
            score += WEIGHTS['three']
        elif count == 2:
            score += WEIGHTS['two']
        elif count == 1:
            score += WEIGHTS['one']
    
    # 撤销模拟落子
    board[row][col] = 0
    
    # 模拟对方落子，检查需要阻挡的情况
    board[row][col] = opponent
    if check_win(row, col):
        score += WEIGHTS['block_win']
    
    for dr, dc in directions:
        # 检查对方的连珠
        count = 1
        # 正方向
        r, c = row + dr, col + dc
        while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and board[r][c] == opponent:
            count += 1
            r += dr
            c += dc
        # 反方向
        r, c = row - dr, col - dc
        while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and board[r][c] == opponent:
            count += 1
            r -= dr
            c -= dc
        
        # 根据连珠数量加分（阻挡）
        if count == 4:
            score += WEIGHTS['block_four']
        elif count == 3:
            score += WEIGHTS['block_three']
        elif count == 2:
            score += WEIGHTS['block_two']
        elif count == 1:
            score += WEIGHTS['block_one']
    
    # 撤销模拟落子
    board[row][col] = 0
    
    # 中心位置加成
    center_row, center_col = BOARD_SIZE // 2, BOARD_SIZE // 2
    distance = abs(row - center_row) + abs(col - center_col)
    score += (10 - distance) * 2
    
    return score

def ai_move():
    """AI落子"""
    best_score = -float('inf')
    best_move = None
    
    # 遍历棋盘上的所有空位
    for row in range(BOARD_SIZE):
        for col in range(BOARD_SIZE):
            if board[row][col] == 0:
                # 评估该位置的价值
                score = evaluate_position(row, col, current_player)
                if score > best_score:
                    best_score = score
                    best_move = (row, col)
    
    # 执行最佳落子
    if best_move:
        row, col = best_move
        board[row][col] = current_player
        return (row, col)
    return None

def show_winner():
    """显示获胜信息"""
    # 尝试使用系统中文字体
    try:
        font = pygame.font.SysFont(["SimHei", "Microsoft YaHei", "Arial"], 60)
    except:
        font = pygame.font.Font(None, 60)
    text = f"{'黑方' if winner == 1 else '白方'}获胜！"
    text_surface = font.render(text, True, (255, 0, 0))
    text_rect = text_surface.get_rect(center=(SCREEN_SIZE//2, SCREEN_SIZE//2))
    screen.blit(text_surface, text_rect)
    
    # 显示重新开始提示
    try:
        font_small = pygame.font.SysFont(["SimHei", "Microsoft YaHei", "Arial"], 30)
    except:
        font_small = pygame.font.Font(None, 30)
    restart_text = "按R键重新开始 | 按Q键退出"
    restart_surface = font_small.render(restart_text, True, BLACK)
    restart_rect = restart_surface.get_rect(center=(SCREEN_SIZE//2, SCREEN_SIZE//2 + 50))
    screen.blit(restart_surface, restart_rect)

def show_mode_selection():
    """显示游戏模式选择界面"""
    # 填充背景色
    screen.fill(BG_COLOR)
    
    # 显示标题
    try:
        title_font = pygame.font.SysFont(["SimHei", "Microsoft YaHei", "Arial"], 48)
    except:
        title_font = pygame.font.Font(None, 48)
    title_text = "五子棋"
    title_surface = title_font.render(title_text, True, BLACK)
    title_rect = title_surface.get_rect(center=(SCREEN_SIZE//2, SCREEN_SIZE//2 - 100))
    screen.blit(title_surface, title_rect)
    
    # 显示模式选择
    try:
        font = pygame.font.SysFont(["SimHei", "Microsoft YaHei", "Arial"], 36)
    except:
        font = pygame.font.Font(None, 36)
    
    # 双人对战选项
    pvp_text = "1. 双人对战"
    pvp_surface = font.render(pvp_text, True, BLACK)
    pvp_rect = pvp_surface.get_rect(center=(SCREEN_SIZE//2, SCREEN_SIZE//2))
    screen.blit(pvp_surface, pvp_rect)
    
    # AI对战选项
    ai_text = "2. AI对战"
    ai_surface = font.render(ai_text, True, BLACK)
    ai_rect = ai_surface.get_rect(center=(SCREEN_SIZE//2, SCREEN_SIZE//2 + 60))
    screen.blit(ai_surface, ai_rect)
    
    # 显示操作提示
    try:
        hint_font = pygame.font.SysFont(["SimHei", "Microsoft YaHei", "Arial"], 24)
    except:
        hint_font = pygame.font.Font(None, 24)
    hint_text = "请按数字键选择游戏模式"
    hint_surface = hint_font.render(hint_text, True, BLACK)
    hint_rect = hint_surface.get_rect(center=(SCREEN_SIZE//2, SCREEN_SIZE//2 + 120))
    screen.blit(hint_surface, hint_rect)
    
    pygame.display.update()

def restart_game():
    """重新开始游戏"""
    global board, current_player, game_over, winner, mode_selected
    board = [[0 for _ in range(BOARD_SIZE)] for _ in range(BOARD_SIZE)]
    current_player = 1
    game_over = False
    winner = None
    mode_selected = False  # 重新开始时需要重新选择模式

def main():
    """游戏主循环"""
    global current_player, game_over, winner, game_mode, mode_selected
    
    clock = pygame.time.Clock()
    
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            
            # 键盘事件处理
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:  # 按R重新开始
                    restart_game()
                if event.key == pygame.K_q:  # 按Q退出
                    pygame.quit()
                    sys.exit()
                
                # 游戏模式选择
                if not mode_selected:
                    if event.key == pygame.K_1:
                        game_mode = 0  # 双人对战
                        mode_selected = True
                    elif event.key == pygame.K_2:
                        game_mode = 1  # AI对战
                        mode_selected = True
            
            # 鼠标点击落子（游戏未结束且已选择模式时）
            if event.type == pygame.MOUSEBUTTONDOWN and not game_over and mode_selected:
                pos = pygame.mouse.get_pos()
                board_pos = get_click_pos(pos)
                
                if board_pos:
                    row, col = board_pos
                    # 检查位置是否为空
                    if board[row][col] == 0:
                        board[row][col] = current_player
                        
                        # 检查是否获胜
                        if check_win(row, col):
                            game_over = True
                            winner = current_player
                        else:
                            # 切换玩家
                            current_player = 2 if current_player == 1 else 1
                            
                            # 如果是AI对战模式且当前玩家是AI（白方），则AI落子
                            if game_mode == 1 and current_player == 2:
                                ai_pos = ai_move()
                                if ai_pos:
                                    ai_row, ai_col = ai_pos
                                    # 检查AI落子后是否获胜
                                    if check_win(ai_row, ai_col):
                                        game_over = True
                                        winner = current_player
                                    else:
                                        # 切换回玩家
                                        current_player = 1
        
        # 如果未选择模式，显示模式选择界面
        if not mode_selected:
            show_mode_selection()
        else:
            # 绘制游戏界面
            draw_board()
            
            # 如果游戏结束，显示获胜信息
            if game_over:
                show_winner()
        
        # 更新屏幕显示
        pygame.display.update()
        clock.tick(60)

if __name__ == "__main__":
    main()