chess_engine/promotion.py

# chess_engine/promotion.py

import chess
from typing import List, Optional, Tuple, Dict, Any
from enum import Enum

class PromotionPiece(Enum):
    """Available pieces for pawn promotion"""
    QUEEN = chess.QUEEN
    ROOK = chess.ROOK
    BISHOP = chess.BISHOP
    KNIGHT = chess.KNIGHT

class PromotionMoveHandler:
    """
    Handles pawn promotion move detection, validation, and generation.
    
    This class provides functionality to:
    - Detect when a pawn move would result in promotion
    - Validate promotion moves
    - Generate all possible promotion moves for a given pawn
    - Create UCI notation moves with promotion suffixes
    """
    
    # Mapping from piece names to chess.PieceType
    PIECE_NAME_MAP = {
        'queen': chess.QUEEN,
        'rook': chess.ROOK,
        'bishop': chess.BISHOP,
        'knight': chess.KNIGHT,
        'q': chess.QUEEN,
        'r': chess.ROOK,
        'b': chess.BISHOP,
        'n': chess.KNIGHT
    }
    
    # Mapping from chess.PieceType to UCI notation
    UCI_PIECE_MAP = {
        chess.QUEEN: 'q',
        chess.ROOK: 'r',
        chess.BISHOP: 'b',
        chess.KNIGHT: 'n'
    }
    
    @staticmethod
    def is_promotion_move(board: chess.Board, from_square: str, to_square: str) -> bool:
        """
        Detect if a move from one square to another would result in pawn promotion.
        
        Args:
            board: Current chess board state
            from_square: Starting square in algebraic notation (e.g., 'e7')
            to_square: Destination square in algebraic notation (e.g., 'e8')
            
        Returns:
            True if the move is a pawn promotion, False otherwise
        """
        try:
            from_idx = chess.parse_square(from_square)
            to_idx = chess.parse_square(to_square)
            
            # Get the piece at the from square
            piece = board.piece_at(from_idx)
            
            # Must be a pawn
            if piece is None or piece.piece_type != chess.PAWN:
                return False
            
            # Check if pawn is moving to the back rank
            to_rank = chess.square_rank(to_idx)
            
            # White pawns promote on 8th rank (rank 7 in 0-indexed), black on 1st rank (rank 0)
            if piece.color == chess.WHITE and to_rank == 7:
                return True
            elif piece.color == chess.BLACK and to_rank == 0:
                return True
                
            return False
            
        except ValueError:
            return False
    
    @staticmethod
    def can_pawn_promote(board: chess.Board, square: str) -> bool:
        """
        Check if a pawn at the given square can potentially promote.
        
        Args:
            board: Current chess board state
            square: Square in algebraic notation (e.g., 'e7')
            
        Returns:
            True if the pawn can promote in the next move, False otherwise
        """
        try:
            square_idx = chess.parse_square(square)
            piece = board.piece_at(square_idx)
            
            # Must be a pawn
            if piece is None or piece.piece_type != chess.PAWN:
                return False
            
            rank = chess.square_rank(square_idx)
            
            # White pawns on 7th rank (rank 6 in 0-indexed) can promote
            # Black pawns on 2nd rank (rank 1 in 0-indexed) can promote
            if piece.color == chess.WHITE and rank == 6:
                return True
            elif piece.color == chess.BLACK and rank == 1:
                return True
                
            return False
            
        except ValueError:
            return False
    
    @staticmethod
    def get_promotion_moves(board: chess.Board, square: str) -> List[str]:
        """
        Get all possible promotion moves for a pawn at the given square.
        
        Args:
            board: Current chess board state
            square: Square in algebraic notation (e.g., 'e7')
            
        Returns:
            List of promotion moves in UCI notation (e.g., ['e7e8q', 'e7e8r', 'e7e8b', 'e7e8n'])
        """
        try:
            from_idx = chess.parse_square(square)
            piece = board.piece_at(from_idx)
            
            # Must be a pawn that can promote
            if not PromotionMoveHandler.can_pawn_promote(board, square):
                return []
            
            promotion_moves = []
            
            # Check all legal moves from this square that are promotions
            for move in board.legal_moves:
                if move.from_square == from_idx and move.promotion is not None:
                    promotion_moves.append(move.uci())
            
            return promotion_moves
            
        except ValueError:
            return []
    
    @staticmethod
    def validate_promotion_move(board: chess.Board, move_str: str) -> Tuple[bool, Optional[str]]:
        """
        Validate a promotion move in UCI notation.
        
        Args:
            board: Current chess board state
            move_str: Move in UCI notation (e.g., 'e7e8q')
            
        Returns:
            Tuple of (is_valid, error_message)
        """
        try:
            # Parse the move
            if len(move_str) < 5:
                return False, "Promotion move must include promotion piece (e.g., 'e7e8q')"
            
            from_square = move_str[:2]
            to_square = move_str[2:4]
            promotion_piece = move_str[4:].lower()
            
            # Validate squares
            try:
                from_idx = chess.parse_square(from_square)
                to_idx = chess.parse_square(to_square)
            except ValueError:
                return False, "Invalid square notation"
            
            # Validate promotion piece
            if promotion_piece not in PromotionMoveHandler.UCI_PIECE_MAP.values():
                return False, f"Invalid promotion piece '{promotion_piece}'. Must be one of: q, r, b, n"
            
            # Check if this is actually a promotion move
            if not PromotionMoveHandler.is_promotion_move(board, from_square, to_square):
                return False, "Move is not a valid promotion move"
            
            # Check if the move is legal
            try:
                move = chess.Move.from_uci(move_str)
                if move not in board.legal_moves:
                    return False, "Move is not legal in current position"
            except ValueError:
                return False, "Invalid UCI move format"
            
            return True, None
            
        except Exception as e:
            return False, f"Error validating promotion move: {str(e)}"
    
    @staticmethod
    def create_promotion_move(from_square: str, to_square: str, promotion_piece: str) -> str:
        """
        Create a UCI promotion move string.
        
        Args:
            from_square: Starting square (e.g., 'e7')
            to_square: Destination square (e.g., 'e8')
            promotion_piece: Piece to promote to ('queen', 'rook', 'bishop', 'knight', or 'q', 'r', 'b', 'n')
            
        Returns:
            UCI move string (e.g., 'e7e8q')
        """
        # Normalize promotion piece to UCI format
        piece_lower = promotion_piece.lower()
        if piece_lower in PromotionMoveHandler.PIECE_NAME_MAP:
            piece_type = PromotionMoveHandler.PIECE_NAME_MAP[piece_lower]
            uci_piece = PromotionMoveHandler.UCI_PIECE_MAP[piece_type]
        else:
            # Assume it's already in UCI format
            uci_piece = piece_lower
        
        return f"{from_square}{to_square}{uci_piece}"
    
    @staticmethod
    def get_available_promotions() -> List[Dict[str, Any]]:
        """
        Get list of available promotion pieces with their details.
        
        Returns:
            List of dictionaries containing promotion piece information
        """
        return [
            {
                'type': 'queen',
                'symbol': 'Q',
                'uci': 'q',
                'name': 'Queen',
                'value': 9
            },
            {
                'type': 'rook',
                'symbol': 'R',
                'uci': 'r',
                'name': 'Rook',
                'value': 5
            },
            {
                'type': 'bishop',
                'symbol': 'B',
                'uci': 'b',
                'name': 'Bishop',
                'value': 3
            },
            {
                'type': 'knight',
                'symbol': 'N',
                'uci': 'n',
                'name': 'Knight',
                'value': 3
            }
        ]
    
    @staticmethod
    def parse_promotion_move(move_str: str) -> Optional[Dict[str, str]]:
        """
        Parse a promotion move string and extract components.
        
        Args:
            move_str: Move in UCI notation (e.g., 'e7e8q')
            
        Returns:
            Dictionary with move components or None if invalid
        """
        try:
            if len(move_str) < 5:
                return None
            
            from_square = move_str[:2]
            to_square = move_str[2:4]
            promotion_piece = move_str[4:].lower()
            
            # Validate components
            chess.parse_square(from_square)  # Will raise ValueError if invalid
            chess.parse_square(to_square)   # Will raise ValueError if invalid
            
            if promotion_piece not in PromotionMoveHandler.UCI_PIECE_MAP.values():
                return None
            
            # Convert UCI piece to full name
            piece_name_map = {v: k for k, v in PromotionMoveHandler.UCI_PIECE_MAP.items()}
            piece_type = piece_name_map[promotion_piece]
            piece_names = {
                chess.QUEEN: 'queen',
                chess.ROOK: 'rook',
                chess.BISHOP: 'bishop',
                chess.KNIGHT: 'knight'
            }
            
            return {
                'from_square': from_square,
                'to_square': to_square,
                'promotion_piece_uci': promotion_piece,
                'promotion_piece_name': piece_names[piece_type],
                'full_move': move_str
            }
            
        except ValueError:
            return None