import torch


def get_board_tensor(board_str, black_moves):
    """ function to move from FEN representation to 12x8x8 tensor
    Note: The rows and cols in the tensor correspond to ranks and files
          in the same order, i.e. first row is Rank 1, first col is File A.
          Also, the color to move next occupies the first 6 channels."""

    p_to_int = {"k": 1, "q": 2, "b": 3, "n": 4, "r": 5, "p": 6,
                "K": 7, "Q": 8, "B": 9, "N": 10, "R": 11, "P": 12}
    new_board = torch.zeros(8, 8)
    rank = 7
    file = 0
    for p in board_str:
        if p == "/":
            rank -= 1
            file = 0
        elif not p.isdigit():
            new_board[rank, file] = (p_to_int[p])
            file += 1
        else:
            new_board[rank, file:file+int(p)] = 0
            file += int(p)

    board_tensor = torch.zeros(12, 8, 8)
    pieces = "kqbnrpKQBNRP" if black_moves else "KQBNRPkqbnrp"
    for p_i, p in enumerate(pieces):
        board_tensor[p_i] = new_board == p_to_int[p]

    return board_tensor