game_mechanics.py

"""DO NOT EDIT THIS FILE!

This is a set of functions written by Delta to
 be used by you during the challenge.

Several functions below may be useful to you and can be imported,
these are clearly marked.
"""
import pickle
import random
from pathlib import Path
from time import sleep
from typing import Callable, Dict, Optional, Tuple

import numpy as np
import pygame.gfxdraw


def choose_move_randomly(state: np.ndarray) -> int:
    """Chooses a move randomly from available moves given the state."""
    return random.choice([col for col in range(8) if not is_column_full(state, col)])


def play_connect_4_game(
    your_choose_move: Callable[[np.ndarray], int],
    opponent_choose_move: Callable[[np.ndarray], int],
    game_speed_multiplier: float = 1,
    render: bool = False,
    verbose: bool = False,
) -> int:
    """Play a game where moves are chosen by `your_choose_move()` and `opponent_choose_move()`. Who
    goes first is chosen at random. You can render the game by setting `render=True`.

    Args:
        your_choose_move: function that chooses move (takes state as input)
        opponent_choose_move: function that picks your opponent's next move
        game_speed_multiplier: multiplies the speed of the game. High == fast
        render: whether to render the game using pygame or not
        verbose: whether to print board states to console. Useful for debugging

    Returns: total_return, which is the sum of return from the game
    """
    total_return = 0
    game = Connect4Env(opponent_choose_move)
    state, reward, done, info = game.reset(verbose)
    if render:
        game.render()
    sleep(1 / game_speed_multiplier)

    while not done:
        action = your_choose_move(state)
        state, reward, done, info = game.step(action, verbose)
        if render:
            game.render()
        total_return += reward
        sleep(1 / game_speed_multiplier)

    return total_return


def get_empty_board(num_rows: int = 6, num_cols: int = 8) -> np.ndarray:
    return np.zeros((num_rows, num_cols))


# Loading and saving dictionaries as files:


def get_dict_filepath(team_name: str) -> Path:
    return Path(__file__).parent.resolve() / f"dict_{team_name}.pkl"


def save_dictionary(dict_to_save: Dict, team_name: str) -> None:
    """Use this to save your value function dictionary."""
    with get_dict_filepath(team_name).open("wb") as f:
        pickle.dump(dict_to_save, f)


def load_dictionary(team_name: str) -> Dict:
    """Loads in a previously-saved value function dictionary."""
    with get_dict_filepath(team_name).open("rb") as f:
        return pickle.load(f)


########## POTENTIALLY USEFUL FEATURES ############################


def has_won(board: np.ndarray, column_index: int) -> bool:
    """Checks if a player has won based on the most recently-added piece on the board.

    Args:
        board: The board to check for a win
        column_index: The column whose top piece you should
                        check for 4-in-a-row

    Returns: True if game won, False otherwise
    """
    row_idx = get_top_piece_row_index(board, column_index)
    assert row_idx is not None, f"Column {column_index} is empty in board shown below.\n\n{board}"
    return get_piece_longest_line_length(board, (row_idx, column_index)) >= 4


def is_column_full(board: np.ndarray, column_idx: int) -> bool:
    """Checks if a board column is full of pieces.

    Args:
        board: The board to check
        column_idx: The column to check

    Returns:
        True if the column is full, False otherwise
    """
    return len(np.where(board[:, column_idx] == 0)[0]) == 0


def get_top_piece_row_index(board: np.ndarray, column_idx: int) -> Optional[int]:
    """Gets the row index of the top piece in a specified column."""
    # Find the empty spaces in the column
    for count, element in enumerate(board[:, column_idx]):
        if element != 0:
            # Find lowest row index in the column
            return count
    # If empty, return None
    return None


def place_piece(board: np.ndarray, column_idx: int, player: int = 1) -> Tuple[np.ndarray, int]:
    """Place a piece from a player on the board. This falls down to the lowest available space in
    the column.

    Args:
        board: board to place the piece on (np array)
        column_idx: column to place the piece in (0 -> 7)
        player: player to place the piece for (1 or -1)

    Returns:
        Tuple of (board, row_index) where the board is updated
         and the row_index is the row index of the added piece.
    """
    assert player in {
        1,
        -1,
    }, f"Invalid player: {player}. Player must be 1 or -1"
    # Find the empty spaces in the column
    top_piece_row_index = get_top_piece_row_index(board, column_idx)
    assert (
        top_piece_row_index != 0
    ), f"Invalid move! Attempted to place a piece in column {column_idx}, but it's full!"
    # Set item in the board to that player's number
    row_idx = top_piece_row_index - 1 if top_piece_row_index else board.shape[0] - 1
    board[row_idx, column_idx] = player
    return board, row_idx


def get_piece_longest_line_length(board: np.ndarray, piece_location: Tuple[int, int]) -> int:
    """Get the length of the longest line of pieces that a piece is in.

    Args:
        board: The board to check
        piece_location: The location of the piece to check (row, col)

    Returns: The length of the longest line of pieces through this piece
    """
    player = board[piece_location]
    directions = [
        [0, 1],
        [1, 1],
        [1, 0],
        [1, -1],
    ]

    direction_row_lengths = []
    # Try all directions
    for direction in directions:
        # We're looking for the longest line through this piece,
        #  start at the piece with a line of 1
        num_in_a_row = 1

        # Try spaces in positive direction
        steps_in_positive_dir = 1

        # Take steps in positive direction until we hit a space not filled by this player's piece
        row = piece_location[0] + steps_in_positive_dir * direction[0]
        col = piece_location[1] + steps_in_positive_dir * direction[1]
        while 0 <= row < board.shape[0] and 0 <= col < board.shape[1] and board[row, col] == player:
            num_in_a_row += 1
            steps_in_positive_dir += 1
            row = piece_location[0] + steps_in_positive_dir * direction[0]
            col = piece_location[1] + steps_in_positive_dir * direction[1]

        # Try spaces in negative direction
        steps_in_negative_dir = 1

        row = piece_location[0] - steps_in_negative_dir * direction[0]
        col = piece_location[1] - steps_in_negative_dir * direction[1]
        while (
            row in range(board.shape[0])
            and col in range(board.shape[1])
            and board[row, col] == player
        ):
            num_in_a_row += 1
            steps_in_negative_dir += 1
            row = piece_location[0] - steps_in_negative_dir * direction[0]
            col = piece_location[1] - steps_in_negative_dir * direction[1]

        # Add all row lengths in each direction to the list
        direction_row_lengths.append(num_in_a_row)

    # Return the longest row length in all directions for this piece
    return max(direction_row_lengths)


### THESE FUNCTIONS ARE LESS USEFUL ###


def board_full(board: np.ndarray) -> bool:
    return np.all(board != 0)


class Connect4Env:
    N_ROWS = 6
    N_COLS = 8

    # Constants for rendering
    DISC_SIZE_RATIO = 0.8
    SQUARE_SIZE = 60

    BLUE_COLOR = (23, 93, 222)
    BACKGROUND_COLOR = (19, 72, 162)
    BLACK_COLOR = (0, 0, 0)
    WHITE_COLOR = (255, 255, 255)
    YELLOW_COLOR = (255, 240, 0)
    RED_COLOR = (255, 0, 0)

    def __init__(
        self,
        opponent_choose_move: Callable[[np.ndarray], int] = choose_move_randomly,
    ):
        self._board_visualizer = np.vectorize(lambda x: "X" if x == 1 else "O" if x == -1 else " ")
        self._opponent_choose_move = opponent_choose_move
        self._screen = None
        self.reset()

    def reset(self, verbose=False):
        """Resets game & takes 1st opponent move if they are chosen to go first."""
        self._board = get_empty_board()
        self._player = random.choice([-1, 1])
        self.done = False
        self.winner = None
        if verbose:
            print(f"Starting game. Player {self._player} has first move\n", self)

        reward = 0

        if self._player == -1:
            # Negative sign is because both players should see themselves as player 1
            opponent_action = self._opponent_choose_move(-self._board)
            reward = -self._step(opponent_action, verbose)

        return self._board, reward, self.done, self.info

    def __repr__(self):
        return str(self._board_visualizer(self._board)) + "\n"

    @property
    def info(self) -> Dict[str, int]:
        return {"player_to_take_next_move": self._player, "winner": self.winner}

    def _step(self, col: int, verbose: bool = False) -> int:
        """Takes 1 turn, internal to this class.

        Do not call
        """
        assert not self.done, "Game is over, call .reset() to start a new game"
        assert (
            isinstance(col, int) and 0 <= col < self._board.shape[1]
        ), f"Col should be an int between 0 and {self._board.shape[1]}"
        assert not is_column_full(self._board, col), "You can't place a counter in a full column!"

        self._board, row = place_piece(self._board, col, self._player)

        # Check for game completion
        won = has_won(self._board, col)
        board_is_full = board_full(self._board)

        reward = 1 if won else 0
        self.winner = self._player if won else None
        self.done = won or board_is_full

        if verbose:
            print(f"Player {self._player} places counter at row {row}, column {col}")
            print(self)
            if won:
                print(f"Player {self._player} has won!\n")
            elif board_is_full:
                print("Board full. It's a tie!")

        # Change self.player only when game isn't over
        self._player *= -1 if not self.done else 1

        return reward

    def step(self, col: int, verbose: bool = False) -> Tuple[np.ndarray, int, bool, Dict[str, int]]:
        """Called by user - takes 2 turns, yours and your opponent's"""
        reward = self._step(col, verbose)

        if not self.done:
            # Negative sign is because both players should see themselves as player 1
            opponent_action = self._opponent_choose_move(-self._board)
            opponent_reward = self._step(opponent_action, verbose)
            # Negative sign is because the opponent's victory is your loss
            reward -= opponent_reward

        return self._board, reward, self.done, self.info

    def __del__(self):
        """Destructor, quit pygame if game over."""
        if self._screen is not None:
            pygame.quit()

    def render(self) -> None:
        """Renders game in pygame."""
        if self._screen is None:
            pygame.init()
            self._screen = pygame.display.set_mode(
                (self.SQUARE_SIZE * self.N_COLS, self.SQUARE_SIZE * self.N_ROWS)
            )

        # Draw background of the board
        pygame.gfxdraw.box(
            self._screen,
            pygame.Rect(
                0,
                0,
                self.N_COLS * self.SQUARE_SIZE,
                self.N_ROWS * self.SQUARE_SIZE,
            ),
            self.BLUE_COLOR,
        )

        # Draw the circles - either as spaces if filled or
        for r in range(self.N_ROWS):
            for c in range(self.N_COLS):
                space = self._board[r, c]
                colour = (
                    self.RED_COLOR
                    if space == 1
                    else self.YELLOW_COLOR
                    if space == -1
                    else self.BACKGROUND_COLOR
                )

                # Anti-aliased circle drawing
                pygame.gfxdraw.aacircle(
                    self._screen,
                    c * self.SQUARE_SIZE + self.SQUARE_SIZE // 2,
                    r * self.SQUARE_SIZE + self.SQUARE_SIZE // 2,
                    int(self.DISC_SIZE_RATIO * self.SQUARE_SIZE / 2),
                    colour,
                )

                pygame.gfxdraw.filled_circle(
                    self._screen,
                    c * self.SQUARE_SIZE + self.SQUARE_SIZE // 2,
                    r * self.SQUARE_SIZE + self.SQUARE_SIZE // 2,
                    int(self.DISC_SIZE_RATIO * self.SQUARE_SIZE / 2),
                    colour,
                )
        pygame.display.update()