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, myisupper):
        class_to_letter = {
            "Bee": "Q",
            "Beetle": "B",
            "Spider": "S",
            "Grasshopper": "G",
            "Ant": "A",
        }
        letter = class_to_letter[self.__class__.__name__]
        if not myisupper:
            letter = letter.lower()
        return [letter, self.p, self.q]
    
    def __repr__(self):
        return ",".join(map(str, self.get_piece_info(True)+[self.team]))
    
    def __str__(self):
        return self.__repr__()

    def validate_jumps(self, board):
        valid_jumps = self.get_valid_jumps(board)
        validated_jumps = []

        for jump in valid_jumps:
            if self.hive_integrity_check(board):
                validated_jumps.append(jump)
        return validated_jumps

    def hive_integrity_check(self, board):
        board_copy = {p: {q: board[p][q] for q in board[p]} for p in board}

        # See if the hive falls apart without the piece (if so, then the move wasnt valid)
        board_copy[self.p][self.q] = None  # Remove the piece from its current position

        # Get all remaining pieces on the board
        remaining_pieces = [
            piece for row in board_copy.values() for piece in row.values() if piece
        ]

        # Start BFS from a random piece
        start_piece = remaining_pieces[0]
        visited = set()
        queue = deque([start_piece])

        while queue:
            current_piece = queue.popleft()
            visited.add(current_piece)

            # Get all neighbors of the current piece
            for neighbor_p, neighbor_q in get_neighbors(
                current_piece.p, current_piece.q
            ):
                neighbor_piece = board_copy.get(neighbor_p, {}).get(neighbor_q)
                if neighbor_piece and neighbor_piece not in visited:
                    queue.append(neighbor_piece)

        # Check if all pieces are connected
        return len(visited) != len(remaining_pieces)

    def get_valid_jumps() -> list:
        raise NotImplementedError


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

    def get_valid_jumps(self, board) -> list:
        valid_moves = []
        for neighbor in get_neighbors(self.p, self.q):
            if neighbor in board and 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):
        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):
        start = (self.p, self.q)
        visited = set()
        queue = deque([(start, 0)])  # Queue of (position, distance)
        valid_moves = []

        while queue:
            current_position, current_distance = queue.popleft()
            if current_distance == 3:
                valid_moves.append(current_position)
                continue

            if current_position in visited:
                continue

            visited.add(current_position)

            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():
            for dp, dq in get_neighbors(0, 0):
                pos = (self.p + dp, self.q + dq)
                while pos in board and board[pos] != "":
                    pos = (pos[0] + dp, pos[1] + dq)
                if (
                    pos in board
                    and board[pos] == ""
                    and (pos[0] != self.p or pos[1] != self.q)
                ):
                    yield pos

        valid_moves = list(generate_moves())
        return valid_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

        # Just a recursive, eh?
        def explore(p, q):
            for dp, dq in get_neighbors(0, 0):
                new_p, new_q = p + dp, q + dq
                next_pos = (new_p, new_q)
                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 next_pos in board and board[next_pos] == "":
                    visited.add(next_pos)
                    yield next_pos
                    yield from explore(next_pos[0], next_pos[1])

        valid_moves = list(explore(self.p, self.q))
        return valid_moves


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.myIsUpper = myIsUpper
        self.algorithmName = "just roll the dice, eh?"
        self.transated_board = dict()
        self.can_place = lambda: sum([v for v in self.myPieces.values()])

    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):
        # mapper dict
        piece_class_mapping = {
            "Q": Bee,
            "B": Beetle,
            "S": Spider,
            "G": Grasshopper,
            "A": Ant,
        }

        translated_board = {p: {} for p in board}
        total_pieces_count = 0
        my_pieces_count = 0

        for p, row in board.items():
            for q, tile_content in row.items():
                if tile_content.isalpha():
                    topmost_piece_letter = tile_content[-1]
                    is_upper = topmost_piece_letter.isupper()
                    piece_class = piece_class_mapping.get(topmost_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] = topmost_piece_letter
                else:
                    translated_board[p][q] = tile_content

        return translated_board, total_pieces_count, my_pieces_count

    @property
    def queen_placed(self):
        return any(isinstance(p, Bee) for p in self.myPieces)
    

    def random_piece(self, pieces):
        return random.choice(pieces) if pieces else None

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

    def get_unplaced_pieces(self):
        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 get_valid_placements(self, translated_board, piece_to_place):
        valid_placements = []

        # If the board is empty, place the piece at the center.
        if not any(row.values() for row in translated_board.values()):
            return [(3, 6)]

        # Iterate over each tile in the board
        for p, row in translated_board.items():
            for q, tile_content in row.items():
                # Check if the tile is empty
                if tile_content == "":
                    # Check if any neighbors are your pieces and none are opponent's pieces
                    neighbors = get_neighbors(p, q)
                    if any(
                        # Check if the neighboring tile is occupied by your team's piece
                        isinstance(translated_board.get(np, {}).get(nq), Piece) and
                        translated_board[np][nq].team == piece_to_place.team
                        for np, nq in neighbors
                    ) and not any(
                        # Check if the neighboring tile is occupied by the opponent's piece
                        isinstance(translated_board.get(np, {}).get(nq), Piece) and
                        translated_board[np][nq].team != piece_to_place.team
                        for np, nq in neighbors
                    ):
                        valid_placements.append((p, q))

        return valid_placements
    
    def mover(self):
        movable_pieces = [
            piece
            for row in self.translated_board.values()
            for piece in row.values()
            if piece and piece.team and piece.validate_jumps(self.translated_board)
        ]
        print(movable_pieces)
        input()
        chosen_piece = self.random_piece(movable_pieces) # -> can be None, usually the cause for an error
        if chosen_piece:
            new_p, new_q = random.choice(
                chosen_piece.validate_jumps(self.translated_board)
            )
            return chosen_piece.get_piece_info(self.myIsUpper) + [new_p, new_q]
    
    def placer(self):
        piece_to_place = self.random_piece(self.get_unplaced_pieces())
        if piece_to_place is None:
            return self.mover()
        if piece_to_place:
            valid_placements = self.get_valid_placements(self.translated_board, piece_to_place)
            new_p, new_q = random.choice(valid_placements)
            return piece_to_place.get_piece_info(self.myIsUpper)[:1] + [
                None,
                None,
                new_p,
                new_q,
            ]

    def update(self):
        self.translated_board, total_pieces_count, my_pieces_count = self.translate_board(self.board)
        return total_pieces_count, my_pieces_count
    def move(self):
        total_pieces_count, my_pieces_count = self.update()

        bee_unplaced = "q" in [k.lower() for k in self.myPieces.keys()] and {k.lower(): v for k,v in self.myPieces.items()}["q"]!=0

        if bee_unplaced and (total_pieces_count > 3 or random.choice([True, False])):
            queen_bee = self.get_piece_class('Q')(None, None, True)
            valid_placements = self.get_valid_placements(self.translated_board, queen_bee)
            if valid_placements:
                new_p, new_q = random.choice(valid_placements)
                return queen_bee.get_piece_info(self.myIsUpper)[:1]+ [None, None, new_p, new_q]

        if total_pieces_count == 0:
            piece_to_place = self.random_piece(self.get_unplaced_pieces())
            return (
                piece_to_place.get_piece_info(self.myIsUpper)[:1] + [None, None, 3, 6]
            )

        elif total_pieces_count == 1:
            for _, row in self.translated_board.items():
                for _, piece in row.items():
                    if piece:
                        adjacent_positions = get_neighbors(piece.p, piece.q)
                        random_position = self.random_piece(adjacent_positions)
                        piece_to_place = self.random_piece(self.get_unplaced_pieces())
                        return (
                            piece_to_place.get_piece_info(self.myIsUpper)[:1]
                            + [None, None, *random_position]
                        )
        elif self.myMove <= 4:
            return self.placer()    
        else:
            move_or_place = random.choice(["move", "place"]) if self.can_place else "move"
            if move_or_place == "place":
                return self.placer()
            else:
                return self.mover()
        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