def get_stockfish_move()

in evals/elsuite/cant_do_that_anymore/scripts/diagonal_dataset_creation.py [0:0]


def get_stockfish_move(stockfish: Stockfish, num_moves_to_consider: int) -> str:
    """
    Gets the next move predicted by stockfish. Gets top n predictions and
    selects randomly weighted by each move's centipawn value
    Filters out bishop promotions, since our variant shouldn't have bishops
    """
    # Get top moves, filter out bad ones
    top_moves = stockfish.get_top_moves(num_moves_to_consider)

    # Filter out bishop promotions
    top_moves = [i for i in top_moves if not i["Move"].endswith("b")]

    # If stockfish considers moves that it knows will lead to mate, only
    # select from these moves
    mates = [i for i in top_moves if i["Mate"] is not None]
    if len(mates) > 0:
        top_moves = mates

    # Ensures centipawn value isn't None
    if all([i["Centipawn"] is None for i in top_moves]):
        for move in top_moves:
            move["Centipawn"] = 1
    else:
        top_moves = [i for i in top_moves if i["Centipawn"] is not None]

    # Makes all centipawns positive
    min_centipawn_value = min([i["Centipawn"] for i in top_moves])
    for move in top_moves:
        move["Centipawn"] += abs(min_centipawn_value)

    # Normalise centipawn to a probability distribution
    centipawn_sum = sum([i["Centipawn"] for i in top_moves])
    for move in top_moves:
        move["prob"] = move["Centipawn"] / centipawn_sum

    # Pick move randomly
    prob = random.uniform(0, 1)
    selected_move = None
    for move in top_moves:
        prob -= move["prob"]
        if prob <= 0:
            selected_move = move["Move"]
            break

    return selected_move