mates.py

"""
Generates 1-ply checkmates for Black in initial 2x2 positions with White's
King-11. May not generate all possible positions.
"""

import cshogi

# First, generates possible SFEN strings. Then, extracts those that are 1-ply
# checkmate. Note that the output may include illegal positions (e.g., Two
# Pawns).

# Possible White pieces in initial positions. Promoted pieces are excluded.
white_pieces: list[str] = ["p", "l", "n", "s", "g", "b", "r"]

# Possible SFEN patterns for a particular square or rank. Pieces with No Moves
# and pieces that give a check are excluded.
square21_pieces: list[str] = ["S", "B", *white_pieces]
rank1_patterns: list[str] = [
    "8k",
    *(f"7{x}k" for x in square21_pieces),
]

square12_pieces: list[str] = ["B", *white_pieces]
square22_pieces: list[str] = ["P", "L", "R", *white_pieces]
rank2_patterns: list[str] = [
    "9",
    *(f"8{x}" for x in square12_pieces),
    *(f"7{x}{y}" for x in square22_pieces for y in ["1", *square12_pieces]),
]

# Assumes Black has all pieces in hand.
pieces_in_hand: dict[str, int] = {"R": 2, "B": 2, "G": 4, "S": 4, "N": 4, "L": 4, "P": 18}


def get_pieces_in_hand_field(board_field: str) -> str:
    """
    Generates the SFEN field for the pieces in hand, with the pieces on the
    given board subtracted.
    """
    field: str = ""

    for piece, count in pieces_in_hand.items():
        remaining: int = count - board_field.upper().count(piece)

        # Too many pieces of one type on the board (e.g., three rooks).
        if remaining < 0:
            raise ValueError

        if remaining > 0:
            if remaining > 1:
                field += str(remaining)

            field += piece

    return field


board: cshogi.Board = cshogi.Board()

# If White is not checkmated after 1 ply, stops the search (i.e., only finds a
# 1-ply checkmate).
dfpn: cshogi.DfPn = cshogi.DfPn(draw_ply=1)

for rank1_pattern in rank1_patterns:
    for rank2_pattern in rank2_patterns:
        board_field: str = f"{rank1_pattern}/{rank2_pattern}/9/9/9/9/9/9/9"

        try:
            pieces_in_hand_field: str = get_pieces_in_hand_field(board_field)
        except ValueError:
            # Skips the position with too many pieces of one type.
            continue

        sfen: str = f"{board_field} b {pieces_in_hand_field} 1"
        board.set_sfen(sfen)

        # Displays the position if a checkmate is found.
        if dfpn.search(board):
            print(board.to_bod())
            print(sfen)
            print()