import gradio as gr
from PIL import Image, ImageDraw
import numpy as np
import math

# Constants
BOARD_SIZE = 9
CELL_SIZE = 50
PIECE_RADIUS = 20
EMPTY = 0
WHITE_SOLDIER = 1
BLACK_SOLDIER = 2
KING = 3
CASTLE = (4, 4)
CAMPS = [
    (0,3), (0,4), (0,5), (1,4),  # Top camp
    (8,3), (8,4), (8,5), (7,4),  # Bottom camp
    (3,0), (4,0), (5,0), (4,1),  # Left camp
    (3,8), (4,8), (5,8), (4,7)   # Right camp
]
ESCAPES = [(i,j) for i in [0,8] for j in range(BOARD_SIZE)] + [(i,j) for j in [0,8] for i in range(BOARD_SIZE) if (i,j) not in CAMPS]
COLORS = {
    'empty': (255, 255, 255),
    'castle': (128, 128, 128),
    'camp': (139, 69, 19),
    'escape': (0, 255, 0),
    'white': (255, 255, 255),
    'black': (0, 0, 0),
    'king': (255, 215, 0),
    'highlight': (255, 255, 0)
}

# Game state class
class TablutState:
    def __init__(self):
        self.board = np.zeros((BOARD_SIZE, BOARD_SIZE), dtype=int)
        self.turn = 'WHITE'
        self.black_in_camps = set(CAMPS)
        self.setup_initial_position()
        self.move_history = []

    def setup_initial_position(self):
        self.board[4, 4] = KING
        white_positions = [(3,4), (4,3), (4,5), (5,4), (2,4), (4,2), (4,6), (6,4)]
        for pos in white_positions:
            self.board[pos] = WHITE_SOLDIER
        for pos in CAMPS:
            self.board[pos] = BLACK_SOLDIER

    def copy(self):
        new_state = TablutState()
        new_state.board = self.board.copy()
        new_state.turn = self.turn
        new_state.black_in_camps = self.black_in_camps.copy()
        new_state.move_history = self.move_history.copy()
        return new_state

# Utility functions
def pos_to_coord(pos):
    row, col = pos
    return f"{chr(ord('A') + col)}{row + 1}"

def is_adjacent_to_castle(pos):
    x, y = pos
    cx, cy = CASTLE
    return (abs(x - cx) == 1 and y == cy) or (abs(y - cy) == 1 and x == cx)

def get_friendly_pieces(turn):
    return [WHITE_SOLDIER, KING] if turn == 'WHITE' else [BLACK_SOLDIER]

def manhattan_distance(pos1, pos2):
    return abs(pos1[0] - pos2[0]) + abs(pos1[1] - pos2[1])

# Game logic functions
def is_valid_move(state, from_pos, to_pos):
    if from_pos == to_pos:
        return False
    piece = state.board[from_pos]
    if state.turn == 'WHITE' and piece not in [WHITE_SOLDIER, KING]:
        return False
    if state.turn == 'BLACK' and piece != BLACK_SOLDIER:
        return False
    if state.board[to_pos] != EMPTY:
        return False
    from_row, from_col = from_pos
    to_row, to_col = to_pos
    if from_row != to_row and from_col != to_col:
        return False
    if from_row == to_row:
        step = 1 if to_col > from_col else -1
        for col in range(from_col + step, to_col, step):
            if state.board[from_row, col] != EMPTY:
                return False
    else:
        step = 1 if to_row > from_row else -1
        for row in range(from_row + step, to_row, step):
            if state.board[row, from_col] != EMPTY:
                return False
    if to_pos == CASTLE and piece != KING:
        return False
    if to_pos in CAMPS:
        if state.turn == 'WHITE' or (state.turn == 'BLACK' and from_pos not in state.black_in_camps):
            return False
    return True

def get_legal_moves(state, from_pos):
    piece = state.board[from_pos]
    if not piece or (state.turn == 'WHITE' and piece not in [WHITE_SOLDIER, KING]) or \
       (state.turn == 'BLACK' and piece != BLACK_SOLDIER):
        return []
    row, col = from_pos
    moves = []
    for r in range(BOARD_SIZE):
        if r != row:
            to_pos = (r, col)
            if is_valid_move(state, from_pos, to_pos):
                moves.append(to_pos)
    for c in range(BOARD_SIZE):
        if c != col:
            to_pos = (row, c)
            if is_valid_move(state, from_pos, to_pos):
                moves.append(to_pos)
    return moves

def is_soldier_captured(state, pos, friendly):
    x, y = pos
    friendly_pieces = get_friendly_pieces(friendly)
    if y > 0 and y < BOARD_SIZE - 1:
        if state.board[x, y-1] in friendly_pieces and state.board[x, y+1] in friendly_pieces:
            return True
    if x > 0 and x < BOARD_SIZE - 1:
        if state.board[x-1, y] in friendly_pieces and state.board[x+1, y] in friendly_pieces:
            return True
    if is_adjacent_to_castle(pos):
        cx, cy = CASTLE
        if x == cx:
            if y < cy and y > 0 and state.board[x, y-1] in friendly_pieces:
                return True
            elif y > cy and y < BOARD_SIZE - 1 and state.board[x, y+1] in friendly_pieces:
                return True
        elif y == cy:
            if x < cx and x > 0 and state.board[x-1, y] in friendly_pieces:
                return True
            elif x > cx and x < BOARD_SIZE - 1 and state.board[x+1, y] in friendly_pieces:
                return True
    if pos in CAMPS:
        return False
    for camp in CAMPS:
        cx, cy = camp
        if (x, y) == (cx + 1, cy) and 0 <= cx < BOARD_SIZE and state.board[cx, cy] in friendly_pieces + [EMPTY]:
            return True
        elif (x, y) == (cx - 1, cy) and 0 <= cx < BOARD_SIZE and state.board[cx, cy] in friendly_pieces + [EMPTY]:
            return True
        elif (x, y) == (cx, cy + 1) and 0 <= cy < BOARD_SIZE and state.board[cx, cy] in friendly_pieces + [EMPTY]:
            return True
        elif (x, y) == (cx, cy - 1) and 0 <= cy < BOARD_SIZE and state.board[cx, cy] in friendly_pieces + [EMPTY]:
            return True
    return False

def is_king_captured(state, pos):
    x, y = pos
    if pos == CASTLE:
        return all(state.board[x + dx, y + dy] == BLACK_SOLDIER for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]
                   if 0 <= x + dx < BOARD_SIZE and 0 <= y + dy < BOARD_SIZE)
    elif is_adjacent_to_castle(pos):
        cx, cy = CASTLE
        dx = cx - x
        dy = cy - y
        free_directions = [d for d in [(-1,0), (1,0), (0,-1), (0,1)] if d != (dx, dy)]
        return all(state.board[x + d[0], y + d[1]] == BLACK_SOLDIER for d in free_directions
                   if 0 <= x + d[0] < BOARD_SIZE and 0 <= y + d[1] < BOARD_SIZE)
    else:
        return is_soldier_captured(state, pos, 'BLACK')

def apply_move(state, from_pos, to_pos):
    new_state = state.copy()
    piece = new_state.board[from_pos]
    new_state.board[to_pos] = piece
    new_state.board[from_pos] = EMPTY
    if new_state.turn == 'BLACK' and from_pos in new_state.black_in_camps and to_pos not in CAMPS:
        new_state.black_in_camps.discard(from_pos)
    captures = []
    opponent = 'BLACK' if new_state.turn == 'WHITE' else 'WHITE'
    for x in range(BOARD_SIZE):
        for y in range(BOARD_SIZE):
            if new_state.board[x, y] == (WHITE_SOLDIER if opponent == 'WHITE' else BLACK_SOLDIER):
                if is_soldier_captured(new_state, (x, y), new_state.turn):
                    captures.append((x, y))
    if opponent == 'WHITE':
        king_pos = find_king_position(new_state)
        if king_pos and is_king_captured(new_state, king_pos):
            captures.append(king_pos)
    for pos in captures:
        new_state.board[pos] = EMPTY
    new_state.turn = 'BLACK' if new_state.turn == 'WHITE' else 'WHITE'
    board_tuple = tuple(new_state.board.flatten())
    new_state.move_history.append(board_tuple)
    return new_state

def find_king_position(state):
    for x in range(BOARD_SIZE):
        for y in range(BOARD_SIZE):
            if state.board[x, y] == KING:
                return (x, y)
    return None

def check_game_status(state):
    king_pos = find_king_position(state)
    if king_pos is None:
        return "BLACK WINS"
    if king_pos in ESCAPES:
        return "WHITE WINS"
    pieces = [(x, y) for x in range(BOARD_SIZE) for y in range(BOARD_SIZE) if
              (state.turn == 'WHITE' and state.board[x, y] in [WHITE_SOLDIER, KING]) or
              (state.turn == 'BLACK' and state.board[x, y] == BLACK_SOLDIER)]
    has_moves = any(get_legal_moves(state, pos) for pos in pieces)
    if not has_moves:
        return "BLACK WINS" if state.turn == 'WHITE' else "WHITE WINS"
    if state.move_history.count(tuple(state.board.flatten())) >= 2:
        return "DRAW"
    return "CONTINUE"

# AI implementation
def evaluate_state(state):
    status = check_game_status(state)
    if status == "WHITE WINS":
        return 1000
    elif status == "BLACK WINS":
        return -1000
    elif status == "DRAW":
        return 0
    king_pos = find_king_position(state)
    if not king_pos:
        return -1000
    min_escape_dist = min(manhattan_distance(king_pos, e) for e in ESCAPES)
    white_pieces = sum(1 for x in range(BOARD_SIZE) for y in range(BOARD_SIZE) if state.board[x, y] == WHITE_SOLDIER)
    black_pieces = sum(1 for x in range(BOARD_SIZE) for y in range(BOARD_SIZE) if state.board[x, y] == BLACK_SOLDIER)
    if state.turn == 'WHITE':
        return -min_escape_dist * 10 + white_pieces * 5 - black_pieces * 3
    else:
        return min_escape_dist * 10 - white_pieces * 3 + black_pieces * 5

def minimax(state, depth, alpha, beta, maximizing_player):
    if depth == 0 or check_game_status(state) != "CONTINUE":
        return evaluate_state(state), None
    if maximizing_player:
        max_eval = -math.inf
        best_move = None
        pieces = [(x, y) for x in range(BOARD_SIZE) for y in range(BOARD_SIZE) if state.board[x, y] == BLACK_SOLDIER]
        for from_pos in pieces:
            for to_pos in get_legal_moves(state, from_pos):
                new_state = apply_move(state, from_pos, to_pos)
                eval_score, _ = minimax(new_state, depth - 1, alpha, beta, False)
                if eval_score > max_eval:
                    max_eval = eval_score
                    best_move = (from_pos, to_pos)
                alpha = max(alpha, eval_score)
                if beta <= alpha:
                    break
        return max_eval, best_move
    else:
        min_eval = math.inf
        best_move = None
        pieces = [(x, y) for x in range(BOARD_SIZE) for y in range(BOARD_SIZE) if state.board[x, y] in [WHITE_SOLDIER, KING]]
        for from_pos in pieces:
            for to_pos in get_legal_moves(state, from_pos):
                new_state = apply_move(state, from_pos, to_pos)
                eval_score, _ = minimax(new_state, depth - 1, alpha, beta, True)
                if eval_score < min_eval:
                    min_eval = eval_score
                    best_move = (from_pos, to_pos)
                beta = min(beta, eval_score)
                if beta <= alpha:
                    break
        return min_eval, best_move

def ai_move(state):
    if state.turn != 'BLACK':
        return state, "Not AI's turn"
    depth = 3
    _, move = minimax(state, depth, -math.inf, math.inf, True)
    if move:
        from_pos, to_pos = move
        new_state = apply_move(state, from_pos, to_pos)
        return new_state, f"AI moved from {pos_to_coord(from_pos)} to {pos_to_coord(to_pos)}"
    return state, "AI has no moves"

# Board visualization
def generate_board_image(state, selected_pos=None):
    img = Image.new('RGB', (BOARD_SIZE * CELL_SIZE, BOARD_SIZE * CELL_SIZE), color=COLORS['empty'])
    draw = ImageDraw.Draw(img)
    for i in range(BOARD_SIZE + 1):
        draw.line([(i * CELL_SIZE, 0), (i * CELL_SIZE, BOARD_SIZE * CELL_SIZE)], fill=(0,0,0), width=1)
        draw.line([(0, i * CELL_SIZE), (BOARD_SIZE * CELL_SIZE, i * CELL_SIZE)], fill=(0,0,0), width=1)
    for x in range(BOARD_SIZE):
        for y in range(BOARD_SIZE):
            pos = (x, y)
            fill = COLORS['empty']
            if pos == CASTLE:
                fill = COLORS['castle']
            elif pos in CAMPS:
                fill = COLORS['camp']
            elif pos in ESCAPES:
                fill = COLORS['escape']
            if pos == selected_pos:
                fill = COLORS['highlight']
            draw.rectangle([(y * CELL_SIZE, x * CELL_SIZE), ((y + 1) * CELL_SIZE, (x + 1) * CELL_SIZE)], fill=fill)
    for x in range(BOARD_SIZE):
        for y in range(BOARD_SIZE):
            piece = state.board[x, y]
            if piece != EMPTY:
                center = (y * CELL_SIZE + CELL_SIZE // 2, x * CELL_SIZE + CELL_SIZE // 2)
                color = COLORS['white'] if piece == WHITE_SOLDIER else COLORS['black'] if piece == BLACK_SOLDIER else COLORS['king']
                draw.ellipse([(center[0] - PIECE_RADIUS, center[1] - PIECE_RADIUS),
                              (center[0] + PIECE_RADIUS, center[1] + PIECE_RADIUS)], fill=color, outline=(0,0,0))
    for i in range(BOARD_SIZE):
        draw.text((5, i * CELL_SIZE + CELL_SIZE // 2 - 5), str(BOARD_SIZE - i), fill=(0,0,0))
        draw.text((i * CELL_SIZE + CELL_SIZE // 2 - 5, BOARD_SIZE * CELL_SIZE - 15), chr(ord('A') + i), fill=(0,0,0))
    return img

# Gradio interface functions
def click_board(state, selected_pos, evt: gr.SelectData):
    if state.turn != 'WHITE':
        return state, "It's the AI's turn", generate_board_image(state), selected_pos
    y = evt.index[0] // CELL_SIZE  # Image coordinates (x, y) map to board (row, col)
    x = evt.index[1] // CELL_SIZE
    pos = (y, x)
    if selected_pos is None:
        if state.board[pos] in [WHITE_SOLDIER, KING]:
            return state, f"Selected {pos_to_coord(pos)}", generate_board_image(state, pos), pos
        else:
            return state, "Select a White piece or King", generate_board_image(state), None
    else:
        if is_valid_move(state, selected_pos, pos):
            new_state = apply_move(state, selected_pos, pos)
            status = check_game_status(new_state)
            if status != "CONTINUE":
                return new_state, status, generate_board_image(new_state), None
            ai_state, ai_message = ai_move(new_state)
            final_status = check_game_status(ai_state)
            message = f"Your move to {pos_to_coord(pos)}. {ai_message}. {final_status if final_status != 'CONTINUE' else ''}"
            return ai_state, message, generate_board_image(ai_state), None
        else:
            return state, "Invalid move", generate_board_image(state), None

def new_game():
    state = TablutState()
    return state, "New game started. Your turn (White).", generate_board_image(state), None

# Define Gradio interface
demo = gr.Blocks(title="Tablut Game")
with demo:
    state = gr.State()
    selected_pos = gr.State(value=None)
    board_image = gr.Image(label="Board", type="pil")
    message_label = gr.Label(label="Message")
    new_game_button = gr.Button("New Game")
    board_image.select(fn=click_board, inputs=[state, selected_pos], outputs=[state, message_label, board_image, selected_pos])
    new_game_button.click(fn=new_game, outputs=[state, message_label, board_image, selected_pos])
    demo.load(fn=new_game, outputs=[state, message_label, board_image, selected_pos])

demo.launch()