import gradio as gr
import numpy as np
import random
import cv2
from enum import Enum
from PIL import Image

class Difficulty(Enum):
    EASY = "Easy"
    MEDIUM = "Medium"
    HARD = "Hard"

class GoGame:
    def __init__(self, size=9):
        self.size = size
        self.board = np.zeros((size, size), dtype=int)
        self.current_player = 1  # 1 for black, -1 for white
        self.pass_count = 0
        self.last_move = None
        self.captured_black = 0
        self.captured_white = 0

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

    def get_group(self, x, y):
        color = self.board[x, y]
        if color == 0:
            return set()
        group = {(x, y)}
        frontier = [(x, y)]
        while frontier:
            current = frontier.pop()
            for nx, ny in self.get_neighbors(*current):
                if self.board[nx, ny] == color and (nx, ny) not in group:
                    group.add((nx, ny))
                    frontier.append((nx, ny))
        return group

    def has_liberties(self, x, y):
        group = self.get_group(x, y)
        for stone_x, stone_y in group:
            for nx, ny in self.get_neighbors(stone_x, stone_y):
                if self.board[nx, ny] == 0:
                    return True
        return False

    def capture_stones(self):
        captured = 0
        for x in range(self.size):
            for y in range(self.size):
                if self.board[x, y] == -self.current_player and not self.has_liberties(x, y):
                    group = self.get_group(x, y)
                    captured += len(group)
                    for gx, gy in group:
                        self.board[gx, gy] = 0
        return captured

    def is_valid_move(self, x, y):
        if self.board[x, y] != 0:
            return False
        # Make a temporary move.
        self.board[x, y] = self.current_player
        valid = self.has_liberties(x, y)
        captures_enemy = self.capture_stones() > 0
        # Undo temporary move.
        self.board[x, y] = 0
        return valid or captures_enemy

    def make_move(self, x, y):
        if not self.is_valid_move(x, y):
            return False
        self.board[x, y] = self.current_player
        captured = self.capture_stones()
        if self.current_player == 1:
            self.captured_white += captured
        else:
            self.captured_black += captured
        self.last_move = (x, y)
        self.current_player *= -1
        return True

    def ai_move(self, difficulty):
        empty_positions = list(zip(*np.where(self.board == 0)))
        if not empty_positions:
            return None
        valid_moves = [pos for pos in empty_positions if self.is_valid_move(*pos)]
        if not valid_moves:
            return None
        if difficulty == Difficulty.EASY:
            move = random.choice(valid_moves)
        elif difficulty == Difficulty.MEDIUM:
            move = self.medium_ai(valid_moves)
        else:
            move = self.hard_ai(valid_moves)
        self.make_move(*move)
        return move

    def medium_ai(self, valid_moves):
        if self.last_move:
            lx, ly = self.last_move
            valid_moves.sort(key=lambda p: abs(p[0] - lx) + abs(p[1] - ly))
            return valid_moves[0]
        return self.center_based_move(valid_moves)

    def center_based_move(self, valid_moves):
        center = self.size // 2
        valid_moves.sort(key=lambda p: abs(p[0] - center) + abs(p[1] - center))
        return valid_moves[0]

    def hard_ai(self, valid_moves):
        for move in valid_moves:
            self.board[move] = self.current_player
            if self.capture_stones() > 0:
                self.board[move] = 0
                return move
            self.board[move] = 0
        opponent = -self.current_player
        self.current_player = opponent
        for move in valid_moves:
            if self.is_valid_move(*move):
                self.board[move] = opponent
                if self.capture_stones() > 0:
                    self.board[move] = 0
                    self.current_player = -opponent
                    return move
                self.board[move] = 0
        self.current_player = -opponent
        return self.medium_ai(valid_moves)

def create_board_image(board, last_move=None, hover_pos=None):
    cell_size = 60
    margin = 40
    total_size = board.shape[0] * cell_size + 2 * margin
    # Create wooden background.
    image = np.full((total_size, total_size, 3), [219, 179, 119], dtype=np.uint8)
    # Draw grid.
    for i in range(board.shape[0]):
        cv2.line(image, (margin + i * cell_size, margin), (margin + i * cell_size, total_size - margin), (0, 0, 0), 2)
        cv2.line(image, (margin, margin + i * cell_size), (total_size - margin, margin + i * cell_size), (0, 0, 0), 2)
    # Draw star points.
    star_points = [(2, 2), (2, 6), (4, 4), (6, 2), (6, 6)]
    for point in star_points:
        cv2.circle(image, (margin + point[1] * cell_size, margin + point[0] * cell_size), 4, (0, 0, 0), -1, cv2.LINE_AA)
    # Draw hover indicator, if any.
    if hover_pos:
        hover_row, hover_col = hover_pos
        center = (margin + hover_col * cell_size, margin + hover_row * cell_size)
        overlay = image.copy()
        cv2.circle(overlay, center, 23, (0, 0, 0), -1, cv2.LINE_AA)
        image = cv2.addWeighted(overlay, 0.5, image, 0.5, 0)
        cv2.circle(image, center, 23, (255, 255, 255), 1, cv2.LINE_AA)
    # Draw stones.
    for i in range(board.shape[0]):
        for j in range(board.shape[1]):
            if board[i, j] != 0:
                center = (margin + j * cell_size, margin + i * cell_size)
                # Shadow effect.
                cv2.circle(image, (center[0] + 2, center[1] + 2), 23, (0, 0, 0), -1, cv2.LINE_AA)
                stone_color = (0, 0, 0) if board[i, j] == 1 else (255, 255, 255)
                cv2.circle(image, center, 23, stone_color, -1, cv2.LINE_AA)
                # 3D highlight.
                if board[i, j] == -1:
                    cv2.circle(image, (center[0] - 5, center[1] - 5), 8, (240, 240, 240), -1, cv2.LINE_AA)
                else:
                    cv2.circle(image, (center[0] - 5, center[1] - 5), 8, (40, 40, 40), -1, cv2.LINE_AA)
    # Highlight last move.
    if last_move:
        row, col = last_move
        center = (margin + col * cell_size, margin + row * cell_size)
        cv2.circle(image, center, 5, (255, 0, 0), -1, cv2.LINE_AA)
    return Image.fromarray(image)

class GradioGoGame:
    def __init__(self):
        self.game = GoGame()
        self.difficulty = Difficulty.EASY
        self.hover_pos = None

    def process_click(self, evt: gr.SelectData):
        cell_size = 60
        margin = 40
        # evt.index returns pixel coordinates as (x, y). Convert to board coordinates.
        col = round((evt.index[0] - margin) / cell_size)
        row = round((evt.index[1] - margin) / cell_size)
        if not (0 <= row < self.game.size and 0 <= col < self.game.size):
            return create_board_image(self.game.board, self.game.last_move), "Invalid click position"
        if not self.game.make_move(row, col):
            return create_board_image(self.game.board, self.game.last_move), "Invalid move (occupied or suicide)"
        # AI Move.
        ai_move = self.game.ai_move(self.difficulty)
        if ai_move is None:
            return create_board_image(self.game.board, self.game.last_move), "Game Over"
        status = f"Black captures: {self.game.captured_white} | White captures: {self.game.captured_black}"
        return create_board_image(self.game.board, self.game.last_move), f"AI moved to: {ai_move}\n{status}"

    def update_hover(self, evt: gr.SelectData):
        cell_size = 60
        margin = 40
        col = round((evt.index[0] - margin) / cell_size)
        row = round((evt.index[1] - margin) / cell_size)
        if 0 <= row < self.game.size and 0 <= col < self.game.size:
            self.hover_pos = (row, col)
        else:
            self.hover_pos = None
        return create_board_image(self.game.board, self.game.last_move, self.hover_pos)

    def reset_game(self):
        self.game = GoGame()
        self.hover_pos = None
        return create_board_image(self.game.board), "Game reset"

def create_interface():
    game = GradioGoGame()
    with gr.Blocks(theme=gr.themes.Soft()) as interface:
        gr.Markdown("""
        # Go Game vs AI  
        Play the ancient game of Go against an AI opponent. Black plays first.
        
        **Rules:**  
        - Click on any intersection to place a stone.  
        - Capture enemy stones by surrounding them.  
        - Stones must have liberties (empty adjacent points) to survive.
        """)
        with gr.Row():
            with gr.Column(scale=2):
                board_output = gr.Image(
                    value=create_board_image(game.game.board),
                    label="Go Board",
                    type="pil",
                    height=600,
                    width=600,
                    interactive=True
                )
                msg_output = gr.Textbox(
                    label="Game Status",
                    value="Click on an intersection to place a stone",
                    lines=2
                )
            with gr.Column(scale=1):
                difficulty = gr.Radio(
                    choices=[d.value for d in Difficulty],
                    value=Difficulty.EASY.value,
                    label="AI Difficulty"
                )
                reset_btn = gr.Button("Reset Game", variant="secondary")
        # Update difficulty when the radio value changes.
        def update_difficulty(value):
            game.difficulty = Difficulty(value)
            return value
        difficulty.change(update_difficulty, inputs=[difficulty], outputs=[difficulty])
        # Use .select() for click events on the interactive image.
        board_output.select(
            game.process_click,
            inputs=[],  # No additional inputs: the event data is passed automatically.
            outputs=[board_output, msg_output]
        )
        # (Removed mousemove event since gr.Image doesn't support it in your version.)
        reset_btn.click(
            game.reset_game,
            outputs=[board_output, msg_output]
        )
    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()