import pygame
import sys
import math
import random
import copy

# --- 初始化 ---
pygame.init()
BOARD_SIZE = 9  # 9路棋盘，适合快速对局
CELL = 50
MARGIN = 50
WIDTH = HEIGHT = CELL * (BOARD_SIZE - 1) + MARGIN * 2
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 围棋 - AI vs Human (9路)")
clock = pygame.time.Clock()

# 颜色
BG_COLOR = (220, 180, 120)
LINE_COLOR = (80, 60, 40)
BLACK_COLOR = (30, 30, 30)
WHITE_COLOR = (240, 240, 240)
HIGHLIGHT_COLOR = (255, 50, 50)
PREVIEW_COLOR = (100, 100, 100, 120)
TEXT_COLOR = (50, 50, 50)
TERRITORY_COLOR = (100, 200, 100, 80)

FONT = pygame.font.SysFont("consolas", 24)
BIG_FONT = pygame.font.SysFont("consolas", 50)

# --- 围棋逻辑 ---

class GoGame:
    def __init__(self):
        self.board = [[0]*BOARD_SIZE for _ in range(BOARD_SIZE)]  # 0=空, 1=黑, 2=白
        self.current_player = 1  # 1=黑, 2=白
        self.history = []  # [(x, y, player, captured)]
        self.ko_point = None  # 打劫禁入点 (x, y)
        self.last_move = None
        self.anim_timer = 0
        self.ai_enabled = True
        self.game_over = False
        self.territory = [[0]*BOARD_SIZE for _ in range(BOARD_SIZE)]  # 1=黑地, 2=白地
        self.score_msg = ""

    def reset(self):
        self.board = [[0]*BOARD_SIZE for _ in range(BOARD_SIZE)]
        self.current_player = 1
        self.history = []
        self.ko_point = None
        self.last_move = None
        self.game_over = False
        self.territory = [[0]*BOARD_SIZE for _ in range(BOARD_SIZE)]
        self.score_msg = ""

    def get_neighbors(self, x, y):
        neighbors = []
        for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < BOARD_SIZE and 0 <= ny < BOARD_SIZE:
                neighbors.append((nx, ny))
        return neighbors

    def get_group(self, x, y):
        """获取连通块"""
        if self.board[y][x] == 0: return set(), set()
        color = self.board[y][x]
        group = set()
        liberties = set()
        stack = [(x, y)]
        visited = set()

        while stack:
            cx, cy = stack.pop()
            if (cx, cy) in visited: continue
            visited.add((cx, cy))
            group.add((cx, cy))

            for nx, ny in self.get_neighbors(cx, cy):
                if self.board[ny][nx] == color and (nx, ny) not in visited:
                    stack.append((nx, ny))
                elif self.board[ny][nx] == 0:
                    liberties.add((nx, ny))

        return group, liberties

    def remove_group(self, group):
        """提子"""
        captured = 0
        for x, y in group:
            self.board[y][x] = 0
            captured += 1
        return captured

    def is_valid_move(self, x, y, player):
        """检查是否合法"""
        if self.board[y][x] != 0: return False
        if self.ko_point == (x, y): return False  # 打劫禁入

        # 模拟落子
        self.board[y][x] = player
        opp = 3 - player

        # 检查是否能提子
        can_capture = False
        for nx, ny in self.get_neighbors(x, y):
            if self.board[ny][nx] == opp:
                group, liberties = self.get_group(nx, ny)
                if len(liberties) == 0:
                    can_capture = True
                    break

        # 检查自杀
        group, liberties = self.get_group(x, y)
        self.board[y][x] = 0  # 恢复

        if len(liberties) == 0 and not can_capture:
            return False  # 禁入点

        return True

    def place(self, x, y):
        """落子"""
        if self.game_over or not self.is_valid_move(x, y, self.current_player):
            return False

        player = self.current_player
        opp = 3 - player
        self.board[y][x] = player
        captured_total = 0

        # 提子
        for nx, ny in self.get_neighbors(x, y):
            if self.board[ny][nx] == opp:
                group, liberties = self.get_group(nx, ny)
                if len(liberties) == 0:
                    captured_total += self.remove_group(group)

        # 打劫检测
        group, liberties = self.get_group(x, y)
        if len(group) == 1 and len(liberties) == 1 and captured_total == 1:
            # 可能是打劫
            lib_pos = list(liberties)[0]
            self.ko_point = lib_pos
        else:
            self.ko_point = None

        self.history.append((x, y, player, captured_total))
        self.last_move = (x, y)
        self.anim_timer = 10
        self.current_player = opp
        return True

    def undo(self):
        """悔棋"""
        if not self.history or self.game_over: return
        steps = 2 if self.ai_enabled and len(self.history) >= 2 else 1
        for _ in range(steps):
            if self.history:
                x, y, player, captured = self.history.pop()
                self.board[y][x] = 0
                self.current_player = player
                # 恢复提子（简化：不恢复具体位置，仅恢复数量）
                # 实际应记录被提子位置
        self.last_move = self.history[-1][:2] if self.history else None
        self.ko_point = None

    def calculate_territory(self):
        """数子法计算领地"""
        visited = [[False]*BOARD_SIZE for _ in range(BOARD_SIZE)]
        self.territory = [[0]*BOARD_SIZE for _ in range(BOARD_SIZE)]
        black_score = 0
        white_score = 0

        for y in range(BOARD_SIZE):
            for x in range(BOARD_SIZE):
                if self.board[y][x] != 0 or visited[y][x]:
                    continue

                # BFS 找空区域
                region = []
                stack = [(x, y)]
                borders = set()

                while stack:
                    cx, cy = stack.pop()
                    if visited[cy][cx]: continue
                    visited[cy][cx] = True
                    region.append((cx, cy))

                    for nx, ny in self.get_neighbors(cx, cy):
                        if self.board[ny][nx] == 0 and not visited[ny][nx]:
                            stack.append((nx, ny))
                        elif self.board[ny][nx] != 0:
                            borders.add(self.board[ny][nx])

                # 判断归属
                if len(borders) == 1:
                    owner = borders.pop()
                    for rx, ry in region:
                        self.territory[ry][rx] = owner
                        if owner == 1: black_score += 1
                        else: white_score += 1

        # 加上棋子数
        for y in range(BOARD_SIZE):
            for x in range(BOARD_SIZE):
                if self.board[y][x] == 1: black_score += 1
                elif self.board[y][x] == 2: white_score += 1

        # 贴目（黑贴6.5）
        black_score -= 6.5
        self.score_msg = f"Black: {black_score:.1f} | White: {white_score:.1f} | {'Black Wins!' if black_score > white_score else 'White Wins!'}"
        self.game_over = True

    def ai_move(self):
        """AI 落子（启发式 + 简单模拟）"""
        if self.game_over or self.current_player != 2: return

        best_score = -1
        best_pos = None

        for y in range(BOARD_SIZE):
            for x in range(BOARD_SIZE):
                if not self.is_valid_move(x, y, 2): continue

                score = self.evaluate_move(x, y, 2)
                if score > best_score:
                    best_score = score
                    best_pos = (x, y)

        if best_pos:
            self.place(*best_pos)

    def evaluate_move(self, x, y, player):
        """评估落子价值"""
        score = 0
        opp = 3 - player

        # 模拟落子
        self.board[y][x] = player

        # 1. 提子价值
        for nx, ny in self.get_neighbors(x, y):
            if self.board[ny][nx] == opp:
                group, liberties = self.get_group(nx, ny)
                if len(liberties) == 0:
                    score += len(group) * 10

        # 2. 自身气数
        group, liberties = self.get_group(x, y)
        score += len(liberties) * 2

        # 3. 靠近边线扣分（9路棋盘）
        if x == 0 or x == BOARD_SIZE-1 or y == 0 or y == BOARD_SIZE-1:
            score -= 3

        # 4. 靠近星位加分
        star_points = [(2,2), (2,6), (6,2), (6,6), (4,4)]
        if (x, y) in star_points:
            score += 5

        # 5. 连接己方棋子
        for nx, ny in self.get_neighbors(x, y):
            if self.board[ny][nx] == player:
                score += 3

        self.board[y][x] = 0  # 恢复
        return score

    def draw(self, surface, mouse_pos):
        surface.fill(BG_COLOR)

        # 棋盘线
        for i in range(BOARD_SIZE):
            pos = MARGIN + i * CELL
            pygame.draw.line(surface, LINE_COLOR, (MARGIN, pos), (WIDTH - MARGIN, pos), 2)
            pygame.draw.line(surface, LINE_COLOR, (pos, MARGIN), (pos, HEIGHT - MARGIN), 2)

        # 星位
        stars = [(2,2), (2,6), (6,2), (6,6), (4,4)] if BOARD_SIZE == 9 else []
        for sx, sy in stars:
            pygame.draw.circle(surface, LINE_COLOR, (MARGIN + sx*CELL, MARGIN + sy*CELL), 4)

        # 领地显示
        if self.game_over:
            for y in range(BOARD_SIZE):
                for x in range(BOARD_SIZE):
                    if self.territory[y][x] == 1:
                        s = pygame.Surface((CELL-4, CELL-4), pygame.SRCALPHA)
                        s.fill((0, 0, 0, 60))
                        surface.blit(s, (MARGIN + x*CELL - (CELL-4)//2, MARGIN + y*CELL - (CELL-4)//2))
                    elif self.territory[y][x] == 2:
                        s = pygame.Surface((CELL-4, CELL-4), pygame.SRCALPHA)
                        s.fill((255, 255, 255, 60))
                        surface.blit(s, (MARGIN + x*CELL - (CELL-4)//2, MARGIN + y*CELL - (CELL-4)//2))

        # 棋子
        for y in range(BOARD_SIZE):
            for x in range(BOARD_SIZE):
                if self.board[y][x]:
                    cx = MARGIN + x * CELL
                    cy = MARGIN + y * CELL
                    color = BLACK_COLOR if self.board[y][x] == 1 else WHITE_COLOR
                    r = CELL // 2 - 2
                    if self.last_move == (x, y) and self.anim_timer > 0:
                        r = int(r * (1 + self.anim_timer * 0.02))
                    pygame.draw.circle(surface, color, (cx, cy), r)
                    if self.board[y][x] == 2:
                        pygame.draw.circle(surface, (180, 180, 180), (cx, cy), r, 2)

        # 最后一步高亮
        if self.last_move and not self.game_over:
            lx, ly = self.last_move
            pygame.draw.circle(surface, HIGHLIGHT_COLOR, (MARGIN + lx*CELL, MARGIN + ly*CELL), 6)

        # 鼠标预览
        if not self.game_over and mouse_pos:
            mx, my = mouse_pos
            gx = round((mx - MARGIN) / CELL)
            gy = round((my - MARGIN) / CELL)
            if 0 <= gx < BOARD_SIZE and 0 <= gy < BOARD_SIZE:
                if self.is_valid_move(gx, gy, self.current_player):
                    s = pygame.Surface((CELL, CELL), pygame.SRCALPHA)
                    color = (*BLACK_COLOR, 120) if self.current_player == 1 else (*WHITE_COLOR, 120)
                    pygame.draw.circle(s, color, (CELL//2, CELL//2), CELL//2 - 2)
                    surface.blit(s, (MARGIN + gx*CELL - CELL//2, MARGIN + gy*CELL - CELL//2))

        # UI
        mode = "AI" if self.ai_enabled else "PVP"
        turn = "Black" if self.current_player == 1 else "White"
        info = f"{mode} | Turn: {turn}"
        if self.game_over:
            info = self.score_msg
        txt = FONT.render(info, True, TEXT_COLOR)
        surface.blit(txt, (MARGIN, HEIGHT - 30))

        hint = FONT.render("Z:Undo | R:Reset | A:AI | P:Score", True, (120, 100, 80))
        surface.blit(hint, (WIDTH - 320, HEIGHT - 30))

# --- 主程序 ---
def main():
    game = GoGame()
    running = True

    while running:
        clock.tick(60)
        mouse_pos = pygame.mouse.get_pos()

        for event in pygame.event.get():
            if event.type == pygame.QUIT: pygame.quit(); sys.exit()
            if event.type == pygame.MOUSEBUTTONDOWN and not game.game_over:
                mx, my = mouse_pos
                gx = round((mx - MARGIN) / CELL)
                gy = round((my - MARGIN) / CELL)
                if 0 <= gx < BOARD_SIZE and 0 <= gy < BOARD_SIZE:
                    if game.place(gx, gy):
                        if game.ai_enabled and not game.game_over:
                            game.ai_move()

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_z: game.undo()
                if event.key == pygame.K_r: game.reset()
                if event.key == pygame.K_a:
                    game.ai_enabled = not game.ai_enabled
                    game.reset()
                if event.key == pygame.K_p and not game.game_over:
                    game.calculate_territory()
                if event.key == pygame.K_ESCAPE: pygame.quit(); sys.exit()

        if game.anim_timer > 0: game.anim_timer -= 1

        game.draw(screen, mouse_pos)
        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()