import pygame
import sys
import math

# --- 1. 初始化与常量设置 ---
pygame.init()

GRID_SIZE = 15
CELL_SIZE = 40
MARGIN = 40
BOARD_PX = (GRID_SIZE - 1) * CELL_SIZE
SCREEN_SIZE = BOARD_PX + 2 * MARGIN

# 颜色定义
C_BG = (220, 179, 92)
C_LINE = (0, 0, 0)
C_BLACK = (15, 15, 20)
C_WHITE = (248, 248, 242)
C_LAST = (220, 60, 60)
C_HOVER = (100, 100, 100, 150)

screen = pygame.display.set_mode((SCREEN_SIZE, SCREEN_SIZE))
pygame.display.set_caption("Pygame 五子棋 (人机对战)")
font = pygame.font.SysFont("simhei", 24)
clock = pygame.time.Clock()

# --- 2. 游戏状态管理 ---
# 0: 空, 1: 玩家(黑), 2: AI(白)
board = [[0 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]
current_player = 1      # 1: 玩家先手
game_over = False
last_move = None
hover_pos = None

# AI 相关设置
AI_PLAYER = 2
HUMAN_PLAYER = 1

# --- 3. 核心逻辑与 AI 算法 ---
def get_grid_pos(mouse_pos):
    x, y = mouse_pos
    col = round((x - MARGIN) / CELL_SIZE)
    row = round((y - MARGIN) / CELL_SIZE)
    if 0 <= row < GRID_SIZE and 0 <= col < GRID_SIZE:
        return row, col
    return None

def check_win(row, col, player):
    directions = [(1, 0), (0, 1), (1, 1), (1, -1)]
    for dr, dc in directions:
        count = 1
        r, c = row + dr, col + dc
        while 0 <= r < GRID_SIZE and 0 <= c < GRID_SIZE and board[r][c] == player:
            count += 1
            r += dr
            c += dc
        r, c = row - dr, col - dc
        while 0 <= r < GRID_SIZE and 0 <= c < GRID_SIZE and board[r][c] == player:
            count += 1
            r -= dr
            c -= dc
        if count >= 5:
            return True
    return False

def evaluate_line(line, player):
    """评估一条线上的棋型得分"""
    opponent = 3 - player
    score = 0
    
    # 简单棋型评估字典
    # 包含己方棋子数，两端开放情况
    if line.count(player) == 4:
        if line.count(0) == 1: score = 100000  # 冲四
        if line.count(0) == 2: score = 1000000 # 活四 (必胜)
    elif line.count(player) == 3:
        if line.count(0) == 2: score = 10000   # 活三
        if line.count(0) == 1: score = 100     # 眠三
    elif line.count(player) == 2:
        if line.count(0) == 3: score = 1000    # 活二
        if line.count(0) == 2: score = 10      # 眠二
        
    # 防守对方棋型
    if line.count(opponent) == 4:
        if line.count(0) == 1: score += 50000  # 堵冲四
        if line.count(0) == 2: score += 800000 # 堵活四
    elif line.count(opponent) == 3:
        if line.count(0) == 2: score += 8000   # 堵活三
        
    return score

def get_ai_move():
    """AI 核心决策函数：遍历空位，计算最高分"""
    best_score = -1
    best_move = None
    
    # 为了效率，只评估周围已有棋子的空位
    candidates = set()
    for r in range(GRID_SIZE):
        for c in range(GRID_SIZE):
            if board[r][c] != 0:
                for dr in range(-2, 3):
                    for dc in range(-2, 3):
                        nr, nc = r + dr, c + dc
                        if 0 <= nr < GRID_SIZE and 0 <= nc < GRID_SIZE and board[nr][nc] == 0:
                            candidates.add((nr, nc))
                            
    # 如果棋盘是空的，下天元
    if not candidates:
        return (GRID_SIZE // 2, GRID_SIZE // 2)

    for r, c in candidates:
        score = 0
        directions = [(1, 0), (0, 1), (1, 1), (1, -1)]
        
        # 模拟落子并评估四个方向
        for dr, dc in directions:
            line = []
            for i in range(-4, 5):
                nr, nc = r + dr * i, c + dc * i
                if 0 <= nr < GRID_SIZE and 0 <= nc < GRID_SIZE:
                    line.append(board[nr][nc])
                else:
                    line.append(-1) # 边界外视为障碍
            score += evaluate_line(line, AI_PLAYER)
            
        if score > best_score:
            best_score = score
            best_move = (r, c)
            
    return best_move

# --- 4. 绘制函数 (与之前相同) ---
def draw_board():
    screen.fill(C_BG)
    for i in range(GRID_SIZE):
        start_h = (MARGIN, MARGIN + i * CELL_SIZE)
        end_h = (MARGIN + BOARD_PX, MARGIN + i * CELL_SIZE)
        start_v = (MARGIN + i * CELL_SIZE, MARGIN)
        end_v = (MARGIN + i * CELL_SIZE, MARGIN + BOARD_PX)
        pygame.draw.line(screen, C_LINE, start_h, end_h, 2)
        pygame.draw.line(screen, C_LINE, start_v, end_v, 2)

def draw_pieces():
    for r in range(GRID_SIZE):
        for c in range(GRID_SIZE):
            if board[r][c] != 0:
                center = (MARGIN + c * CELL_SIZE, MARGIN + r * CELL_SIZE)
                color = C_BLACK if board[r][c] == 1 else C_WHITE
                pygame.draw.circle(screen, color, center, CELL_SIZE // 2 - 2)
    if last_move:
        r, c = last_move
        center = (MARGIN + c * CELL_SIZE, MARGIN + r * CELL_SIZE)
        pygame.draw.circle(screen, C_LAST, center, 5)

def draw_hover():
    if hover_pos and not game_over and current_player == HUMAN_PLAYER:
        r, c = hover_pos
        if board[r][c] == 0:
            center = (MARGIN + c * CELL_SIZE, MARGIN + r * CELL_SIZE)
            hover_surf = pygame.Surface((CELL_SIZE, CELL_SIZE), pygame.SRCALPHA)
            pygame.draw.circle(hover_surf, C_HOVER, (CELL_SIZE//2, CELL_SIZE//2), CELL_SIZE//2 - 2)
            screen.blit(hover_surf, (center[0] - CELL_SIZE//2, center[1] - CELL_SIZE//2))

def draw_status():
    if game_over:
        winner = "黑棋(玩家)" if current_player == AI_PLAYER else "白棋(AI)"
        text = f"游戏结束! {winner} 获胜!"
    else:
        text = "轮到: 玩家(黑棋)" if current_player == HUMAN_PLAYER else "思考中: AI(白棋)..."
    text_surface = font.render(text, True, C_BLACK)
    screen.blit(text_surface, (MARGIN, SCREEN_SIZE - 30))

# --- 5. 游戏主循环 ---
while True:
    clock.tick(60)
    draw_board()
    draw_pieces()
    draw_hover()
    draw_status()
    pygame.display.flip()

    # AI 回合自动落子
    if not game_over and current_player == AI_PLAYER:
        move = get_ai_move()
        if move:
            r, c = move
            board[r][c] = AI_PLAYER
            last_move = (r, c)
            if check_win(r, c, AI_PLAYER):
                game_over = True
            else:
                current_player = HUMAN_PLAYER

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        
        if event.type == pygame.MOUSEMOTION:
            hover_pos = get_grid_pos(event.pos)
            
        # 玩家点击落子
        if event.type == pygame.MOUSEBUTTONDOWN and not game_over and current_player == HUMAN_PLAYER:
            pos = get_grid_pos(event.pos)
            if pos:
                r, c = pos
                if board[r][c] == 0:
                    board[r][c] = HUMAN_PLAYER
                    last_move = (r, c)
                    if check_win(r, c, HUMAN_PLAYER):
                        game_over = True
                    else:
                        current_player = AI_PLAYER