import base as Base
import copy, random, time, math
from PIL import Image, ImageDraw
from collections import deque


# Player template for HIVE --- ALP semestral work
# Vojta Vonasek, 2023


# PUT ALL YOUR IMPLEMENTATION INTO THIS FILE


def get_neighbors(p, q):
    directions = [(0, -1), (1, -1), (1, 0), (0, 1), (-1, 1), (-1, 0)]
    return [(p + dp, q + dq) for dp, dq in directions]


class Piece:
    def get_piece_info(self):
        # Convert the class name to the corresponding letter
        class_to_letter = {
            "Bee": "Q",
            "Beetle": "B",
            "Spider": "S",
            "Grasshopper": "G",
            "Ant": "A",
        }
        letter = class_to_letter[self.__class__.__name__]
        if not self.team:  # If the piece is not on the player's team, make it lowercase
            letter = letter.lower()
        return [letter, self.p, self.q]

    def get_valid_jumps() -> list:
        raise NotImplementedError


class Bee(Piece):
    def __init__(self, board, p, q, team):
        self.board = board
        self.p = p
        self.q = q
        self.team = team  # 'team' should be 'upper' or 'lower' based on the piece case.

    def get_valid_jumps(self) -> list:
        # The bee can move to an adjacent empty space.
        valid_moves = []
        for neighbor in get_neighbors(self.p, self.q):
            # Check if the neighbor is within the bounds of the board and is empty
            if neighbor in self.board and self.board[neighbor] == "":
                valid_moves.append(neighbor)
        return valid_moves


class Beetle(Piece):
    def __init__(self, p, q, team):
        self.p = p
        self.q = q
        self.team = team

    def get_valid_jumps(self, board):
        # The beetle can move to an adjacent space whether it is empty or not.
        valid_moves = get_neighbors(self.p, self.q)
        return valid_moves


class Spider(Piece):
    def __init__(self, p, q, team):
        self.p = p
        self.q = q
        self.team = team

    def get_valid_jumps(self, board):
        # Use BFS to find all tiles exactly three moves away
        start = (self.p, self.q)
        visited = set()  # Keep track of visited tiles
        queue = deque([(start, 0)])  # Queue of (position, distance)
        valid_moves = []

        while queue:
            current_position, current_distance = queue.popleft()
            if current_distance == 3:
                # We have found a tile exactly three moves away
                valid_moves.append(current_position)
                continue  # We don't want to move further from this tile

            if current_position in visited:
                continue  # Already visited this tile

            visited.add(current_position)

            # Get all neighboring positions
            for neighbor in get_neighbors(*current_position):
                if (
                    neighbor in board
                    and board[neighbor] == ""
                    and neighbor not in visited
                ):
                    queue.append((neighbor, current_distance + 1))

        return valid_moves


class Grasshopper(Piece):
    def __init__(self, p, q, team):
        self.p = p
        self.q = q
        self.team = team

    def get_valid_jumps(self, board):
        # Generator function to yield valid moves
        def generate_moves():
            # Check each direction for possible jumps
            for dp, dq in get_neighbors(0, 0):  # Use (0, 0) to get direction vectors
                pos = (self.p + dp, self.q + dq)
                # Continue in the direction until we find a piece to jump over
                while pos in board and board[pos] != "":
                    pos = (pos[0] + dp, pos[1] + dq)
                # If we jumped over at least one piece and landed on an empty space, yield the move
                if (
                    pos in board
                    and board[pos] == ""
                    and (pos[0] != self.p or pos[1] != self.q)
                ):
                    yield pos

        return list(generate_moves())


class Ant(Piece):
    def __init__(self, p, q, team):
        self.p = p
        self.q = q
        self.team = team

    def get_valid_jumps(self, board):
        visited = set()  # Keep track of visited nodes to prevent cycles

        # Recursive function to explore all valid moves
        def explore(p, q):
            for dp, dq in get_neighbors(0, 0):  # Get direction vectors
                new_p, new_q = p + dp, q + dq
                next_pos = (new_p, new_q)
                # Continue in this direction while there are pieces to slide around
                while (
                    next_pos in board
                    and board[next_pos] != ""
                    and next_pos not in visited
                ):
                    visited.add(next_pos)
                    next_pos = (next_pos[0] + dp, next_pos[1] + dq)
                # If we found an empty space, add it as a valid move and explore from there
                if next_pos in board and board[next_pos] == "":
                    visited.add(next_pos)
                    yield next_pos
                    # Explore further moves from this new position
                    yield from explore(next_pos[0], next_pos[1])

        # Start exploring from the Ant's current position
        return list(explore(self.p, self.q))


class Player(Base.Board):
    def __init__(
        self, playerName, myIsUpper, size, myPieces, rivalPieces
    ):  # do not change this line
        Base.Board.__init__(
            self, myIsUpper, size, myPieces, rivalPieces
        )  # do not change this line
        self.playerName = playerName
        self.algorithmName = "just roll the dice, eh?"

    def getAllEmptyCells(self):
        result = []
        for p in self.board:
            for q in self.board[p]:
                if self.isEmpty(p, q, self.board):
                    result.append([p, q])
        return result

    def getAllNonemptyCells(self):
        result = []
        for p in self.board:
            for q in self.board[p]:
                if not self.isEmpty(p, q, self.board):
                    result.append([p, q])
        return result

    def translate_board(self, board):
        # Create a dictionary to map letters to piece classes
        piece_class_mapping = {
            "Q": Bee,
            "B": Beetle,
            "S": Spider,
            "G": Grasshopper,
            "A": Ant,
        }

        translated_board = {p: {} for p in board}
        total_pieces_count = 0  # To count the total number of pieces on the board
        my_pieces_count = 0  # To count the number of your pieces on the board

        for p, row in board.items():
            for q, piece_letter in row.items():
                if (
                    piece_letter.isalpha()
                ):  # Check if it's a letter representing a piece
                    is_upper = piece_letter.isupper()
                    piece_class = piece_class_mapping.get(piece_letter.upper())
                    if piece_class:
                        piece_instance = piece_class(p, q, is_upper == self.myIsUpper)
                        translated_board[p][q] = piece_instance
                        total_pieces_count += 1
                        if is_upper == self.myIsUpper:
                            my_pieces_count += 1
                    else:
                        translated_board[p][q] = piece_letter
                else:
                    translated_board[p][q] = piece_letter

        return translated_board, total_pieces_count, my_pieces_count

    @property
    def queen_placed(self):
        # Check if the queen is placed
        return any(isinstance(p, Bee) for p in self.myPieces)

    def get_piece_class(self, letter):
        # Maps letter to piece class
        return {"Q": Bee, "B": Beetle, "S": Spider, "G": Grasshopper, "A": Ant}.get(
            letter
        )

    def get_unplaced_pieces(self):
        # Return a list of unplaced pieces
        unplaced_pieces = []
        for piece_letter, count in self.myPieces.items():
            for _ in range(count):
                piece_class = self.get_piece_class(piece_letter)
                unplaced_pieces.append(piece_class(None, None, True))
        return unplaced_pieces

    def move(self):
        translated_board, total_pieces_count, _ = self.translate_board(self.board)

        # Randomly choose a piece to place or move
        def choose_random_piece(pieces):
            return random.choice(pieces) if pieces else None

        # If we're the first player and no pieces are on the board, place a random piece at the center
        if total_pieces_count == 0:
            piece_to_place = choose_random_piece(self.get_unplaced_pieces())
            return (
                piece_to_place.get_piece_info()[:1] + [None, None, 3, 6]
                if piece_to_place
                else []
            )

        # If we're the second player, place next to the first player's piece
        elif total_pieces_count == 1:
            for _, row in translated_board.items():
                for _, piece in row.items():
                    if piece:  # Found the first player's piece
                        adjacent_positions = get_neighbors(piece.p, piece.q)
                        random_position = choose_random_piece(adjacent_positions)
                        piece_to_place = choose_random_piece(self.get_unplaced_pieces())
                        return (
                            piece_to_place.get_piece_info()[:1]
                            + [None, None, *random_position]
                            if piece_to_place
                            else []
                        )

        # After the queen is placed or after the 4th turn, randomly choose between moving and placing a piece
        elif self.queen_placed or total_pieces_count >= 8:
            move_or_place = random.choice(["move", "place"])
            if move_or_place == "move":
                movable_pieces = [
                    p
                    for p in translated_board.values()
                    if p.get_valid_jumps(translated_board)
                ]
                chosen_piece = choose_random_piece(movable_pieces)
                if chosen_piece:
                    new_p, new_q = random.choice(
                        chosen_piece.get_valid_jumps(translated_board)
                    )
                    return chosen_piece.get_piece_info() + [new_p, new_q]
            else:
                piece_to_place = choose_random_piece(self.get_unplaced_pieces())
                if piece_to_place:
                    valid_placements = self.get_valid_placements(
                        translated_board, piece_to_place
                    )
                    new_p, new_q = (
                        random.choice(valid_placements)
                        if valid_placements
                        else (None, None)
                    )
                    return piece_to_place.get_piece_info()[:1] + [
                        None,
                        None,
                        new_p,
                        new_q,
                    ]

        # If it's not possible to move or place a piece, pass the turn
        return []


def updatePlayers(move, activePlayer, passivePlayer):
    """write move made by activePlayer player
    this method assumes that all moves are correct, no checking is made
    """
    if len(move) == 0:
        return

    animal, p, q, newp, newq = move
    if p == None and q == None:
        # placing new animal
        activePlayer.myPieces[animal] -= 1
        passivePlayer.rivalPieces = activePlayer.myPieces.copy()
    else:
        # just moving animal
        # delete its old position
        activePlayer.board[p][q] = activePlayer.board[p][q][:-1]
        passivePlayer.board[p][q] = passivePlayer.board[p][q][:-1]

    activePlayer.board[newp][newq] += animal
    passivePlayer.board[newp][newq] += animal


if __name__ == "__main__":
    boardSize = 13
    smallFigures = {
        "q": 1,
        "a": 2,
        "b": 2,
        "s": 2,
        "g": 2,
    }  # key is animal, value is how many is available for placing
    bigFigures = {
        figure.upper(): smallFigures[figure] for figure in smallFigures
    }  # same, but with upper case

    P1 = Player("player1", False, 13, smallFigures, bigFigures)
    P2 = Player("player2", True, 13, bigFigures, smallFigures)

    filename = "begin.png"
    P1.saveImage(filename)

    moveIdx = 0
    while True:
        move = P1.move()
        print("P1 returned", move)
        updatePlayers(move, P1, P2)  # update P1 and P2 according to the move
        filename = "move-{:03d}-player1.png".format(moveIdx)
        P1.saveImage(filename)

        move = P2.move()
        print("P2 returned", move)
        updatePlayers(move, P2, P1)  # update P2 and P1 according to the move
        filename = "move-{:03d}-player2.png".format(moveIdx)
        P1.saveImage(filename)

        moveIdx += 1
        P1.myMove = moveIdx
        P2.myMove = moveIdx

        if moveIdx > 50:
            print("End of the test game")
            break