Python Bit Board Tic-Tac-Toe

# We'll use this package to get random moves for the computer player.
from random import randint

"""
GLOSSARY
- A **line** is a row, column or diagonal.
- A **bit word** is a binary number used to represent data of some sort.
- A bit is said to be **set** when it is equal to 1. For example, in `0100` the bit at position 2 (counting from the right from zero) is set.
- A **mask** is a bit word used to check if certain bits in another number are set. Examples:
    - `n & 0011` is used to check if the first two bits in the number `n` are set,
    - `n2 = n | 0001` returns a number `n2` equal to `n` with its first bit set to 1.
- A **bit board** is a bit word used specifically to represent the board in a board game.

BOARD REPRESENTATION
The board is represented as a single bit word (a bit board).

bits 00-08 = X occupancy
bits 09-17 = O occupancy
bit 18     = active player

CELL INDICES
By convention, indices from 0 to 9 represent the cells on the board in the order below:

6 7 8
3 4 5
0 1 2

EXAMPLE
Given the board below:

O . X
. X .
. . .

the corresponding bit word is: 1 001000000 100010000
100010000 => X has played to cells 4 and 8
001000000 => O has played to cell 6
1         => it is O's turn to move
"""

NB_ROWS = 3
""" Number of rows on the board. """

NB_COLS = NB_ROWS
""" Number of columns per row. """

NB_CELLS = NB_ROWS * NB_COLS
""" Total number of cells on the board. """

PLAYER_COMPUTER = 0
""" The AI player. Always goes first. """

PLAYER_HUMAN = 1
""" The human player. Always goes second. """

PLAYERS_SYMBOLS = ["X", "O"]
""" The characters used to represent each player when printing the board. """

STATE_LOSS = 0
""" A value indicating that the game was lost by the active player. """

STATE_DRAW = 1
""" A value indicating that the game is over and neither player won. """

STATE_ONGOING = 2
""" A value indicating that the game hasn't reached a result yet. """

NB_BITS_FULL_OCCUPANCY = NB_CELLS * 2
""" The number of bits to represent the occupancy of both players. """

NB_BITS_BOARD = NB_BITS_FULL_OCCUPANCY + 1
""" The number of bits in a bit board (including the active player). """

OCCUPANCY_MASK = (1 << NB_CELLS) - 1
""" A mask with 9 set bits. Used to detect if the board is full or single out a player's occupancy. """

BIT_MASKS = [1 << i for i in range(NB_BITS_BOARD)]
""" A list of masks each having a single set bit used for various bit manipulations. """

WIN_MASKS = []  # type: list[int]
""" A list of masks each having 3 set bits used to quickly detect if either player has filled a line. It is initialized later through a function call. """

BIT_INDEX_PLAYER = NB_BITS_FULL_OCCUPANCY
""" The index of the bit representing the active player in a bit board. """


def coordsToIndex(row: int, col: int) -> int:
    """ Convert a row and column to a board index. """
    return row * NB_COLS + col


def initWinMasks() -> None:
    """ The masks should look like this:
        000 000 111,
        000 111 000,
        111 000 000,
        001 001 001,
        010 010 010,
        100 100 100,
        100 010 001,
        001 010 100

        They each represent a line that can be filled to win the game.
    """
    diagonalMask = 0
    antiDiagonalMask = 0

    for row in range(NB_ROWS):
        rowMask = 0
        colMask = 0

        for col in range(NB_COLS):
            rowMask |= BIT_MASKS[coordsToIndex(row, col)]
            colMask |= BIT_MASKS[coordsToIndex(col=row, row=col)]

        WIN_MASKS.append(rowMask)
        WIN_MASKS.append(colMask)
        diagonalMask |= BIT_MASKS[coordsToIndex(row, row)]
        antiDiagonalMask |= BIT_MASKS[coordsToIndex(row, NB_ROWS - row - 1)]

    WIN_MASKS.append(diagonalMask)
    WIN_MASKS.append(antiDiagonalMask)


def getActivePlayer(board: int) -> int:
    return board >> BIT_INDEX_PLAYER


def oppositePlayer(player: int) -> int:
    # Since `player` is always equal to 0 or 1, it can be reversed by toggling it (XOR operator).
    return player ^ 1


def getOccupancy(board: int, player: int) -> int:
    """ Extract the part of a bit board representing the cells occupied by a given player. """
    return board >> (NB_CELLS * player) & OCCUPANCY_MASK


def isCellOccupied(occ: int, cell: int) -> bool:
    bitMask = BIT_MASKS[cell]
    # If the cell wasn't occupied, then `fullOcc & bitMask` would be equal to 0.
    return (occ & bitMask) != 0


def getFreeCells(board: int) -> list[int]:
    """ Get a list of unoccupied cell indices. """
    computerOcc = getOccupancy(board, PLAYER_COMPUTER)
    humanOcc = getOccupancy(board, PLAYER_HUMAN)
    fullOcc = computerOcc | humanOcc
    freeCells = []  # type: list[int]

    for cell in range(NB_CELLS):
        if not isCellOccupied(fullOcc, cell):
            freeCells.append(cell)

    return freeCells


def getState(board: int) -> int:
    """ Check if the game is lost, drawn or still in play. """
    activePlayer = getActivePlayer(board)
    inactivePlayer = oppositePlayer(activePlayer)
    inactiveOcc = getOccupancy(board, inactivePlayer)

    for mask in WIN_MASKS:
        if (inactiveOcc & mask) == mask:
            return STATE_LOSS

    activeOcc = getOccupancy(board, activePlayer)
    fullOcc = activeOcc | inactiveOcc

    if fullOcc == OCCUPANCY_MASK:
        return STATE_DRAW

    return STATE_ONGOING


def promptHumanMove(board: int) -> int:
    """ Ask user to enter a move via the command line. """
    fullOcc = board & ~BIT_MASKS[BIT_INDEX_PLAYER]
    prompt = "Please choose a free cell from the above board: "

    try:
        move = int(input(prompt))

        if move < 0 or move >= NB_CELLS:
            print(f"Invalid move: should be in the range [0;{NB_CELLS}[.")
            # Call function again to query user once more.
            return promptHumanMove(board)

        if isCellOccupied(fullOcc, move):
            print("This cell is occupied.")
            return promptHumanMove(board)

        return move
    except:
        print("Invalid input!")
        return promptHumanMove(board)


def getComputerMove(board: int) -> int:
    freeCells = getFreeCells(board)
    # Pick a random move for the computer to play.
    randomIndex = randint(0, len(freeCells) - 1)
    return freeCells[randomIndex]


def printBoard(board: int) -> None:
    """ Display a visual representation of the board. """

    # Command line colors
    FG_YELLOW = "\033[93m"
    FG_CYAN = "\033[96m"
    CLR_RESET = "\033[0m"

    computerOcc = getOccupancy(board, PLAYER_COMPUTER)
    humanOcc = getOccupancy(board, PLAYER_HUMAN)

    for row in range(NB_ROWS):
        for col in range(NB_COLS):
            # Row is reversed in order to print the first row at the bottom.
            index = coordsToIndex(NB_ROWS - row - 1, col)

            if isCellOccupied(computerOcc, index):
                print(
                    FG_CYAN + PLAYERS_SYMBOLS[PLAYER_COMPUTER] + CLR_RESET, end=" ")
                continue

            if isCellOccupied(humanOcc, index):
                print(
                    FG_YELLOW + PLAYERS_SYMBOLS[PLAYER_HUMAN] + CLR_RESET, end=" ")
                continue

            print(index, end=" ")

        print()


def play() -> None:
    board = 0
    # Counting the number of moves played to avoid an infinite loop.
    nbMovesPlayed = 0

    while nbMovesPlayed <= NB_CELLS:
        activePlayer = getActivePlayer(board)
        inactivePlayer = oppositePlayer(activePlayer)
        state = getState(board)

        if state == STATE_LOSS:
            print(f"{PLAYERS_SYMBOLS[inactivePlayer]} wins!")
            break

        if state == STATE_DRAW:
            print("The game is a draw.")
            break

        if activePlayer == PLAYER_COMPUTER:
            move = getComputerMove(board)
            board |= BIT_MASKS[move]
            print(f"The computer played to cell {move}.")
        elif activePlayer == PLAYER_HUMAN:
            move = promptHumanMove(board)
            board |= BIT_MASKS[move + NB_CELLS]

        print()
        printBoard(board)
        print()

        # Toggle player
        board ^= BIT_MASKS[BIT_INDEX_PLAYER]
        nbMovesPlayed += 1


def main() -> None:
    initWinMasks()
    play()


# Ensure the `main` function is only called when this script is run directly.
if __name__ == "__main__":
    main()