chess_engine/pieces.py

# chess_engine/pieces.py

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

class PieceType(Enum):
    PAWN = chess.PAWN
    KNIGHT = chess.KNIGHT
    BISHOP = chess.BISHOP
    ROOK = chess.ROOK
    QUEEN = chess.QUEEN
    KING = chess.KING

class PieceColor(Enum):
    WHITE = chess.WHITE
    BLACK = chess.BLACK

class ChessPiece:
    """Represents a chess piece with its properties and behaviors"""
    
    # Unicode symbols for pieces
    UNICODE_PIECES = {
        chess.WHITE: {
            chess.PAWN: "♙", chess.KNIGHT: "♘", chess.BISHOP: "♗",
            chess.ROOK: "♖", chess.QUEEN: "♕", chess.KING: "♔"
        },
        chess.BLACK: {
            chess.PAWN: "♟", chess.KNIGHT: "♞", chess.BISHOP: "♝",
            chess.ROOK: "♜", chess.QUEEN: "♛", chess.KING: "♚"
        }
    }
    
    # Piece values for evaluation
    PIECE_VALUES = {
        chess.PAWN: 1,
        chess.KNIGHT: 3,
        chess.BISHOP: 3,
        chess.ROOK: 5,
        chess.QUEEN: 9,
        chess.KING: 0  # King is invaluable
    }
    
    def __init__(self, piece_type: chess.PieceType, color: chess.Color):
        self.piece_type = piece_type
        self.color = color
        self.piece = chess.Piece(piece_type, color)
    
    @property
    def symbol(self) -> str:
        """Get the piece symbol (uppercase for white, lowercase for black)"""
        return self.piece.symbol()
    
    @property
    def unicode_symbol(self) -> str:
        """Get the Unicode symbol for the piece"""
        return self.UNICODE_PIECES[self.color][self.piece_type]
    
    @property
    def value(self) -> int:
        """Get the piece value"""
        return self.PIECE_VALUES[self.piece_type]
    
    @property
    def name(self) -> str:
        """Get the piece name"""
        piece_names = {
            chess.PAWN: "Pawn",
            chess.KNIGHT: "Knight",
            chess.BISHOP: "Bishop",
            chess.ROOK: "Rook",
            chess.QUEEN: "Queen",
            chess.KING: "King"
        }
        return piece_names[self.piece_type]
    
    @property
    def color_name(self) -> str:
        """Get the color name"""
        return "White" if self.color == chess.WHITE else "Black"
    
    def __str__(self) -> str:
        return f"{self.color_name} {self.name}"
    
    def __repr__(self) -> str:
        return f"ChessPiece({self.piece_type}, {self.color})"

class PieceManager:
    """Manages piece-related operations and utilities"""
    
    @staticmethod
    def create_piece(piece_type: str, color: str) -> Optional[ChessPiece]:
        """
        Create a piece from string representations
        
        Args:
            piece_type: 'pawn', 'knight', 'bishop', 'rook', 'queen', 'king'
            color: 'white' or 'black'
            
        Returns:
            ChessPiece instance or None if invalid input
        """
        piece_type_map = {
            'pawn': chess.PAWN,
            'knight': chess.KNIGHT,
            'bishop': chess.BISHOP,
            'rook': chess.ROOK,
            'queen': chess.QUEEN,
            'king': chess.KING
        }
        
        color_map = {
            'white': chess.WHITE,
            'black': chess.BLACK
        }
        
        if piece_type.lower() not in piece_type_map or color.lower() not in color_map:
            return None
            
        return ChessPiece(
            piece_type_map[piece_type.lower()],
            color_map[color.lower()]
        )
    
    @staticmethod
    def get_piece_moves(board: chess.Board, square: str) -> List[str]:
        """
        Get all possible moves for a piece at given square
        
        Args:
            board: Chess board instance
            square: Square in algebraic notation
            
        Returns:
            List of destination squares in algebraic notation
        """
        try:
            square_index = chess.parse_square(square)
            piece = board.piece_at(square_index)
            
            if piece is None:
                return []
            
            moves = []
            for move in board.legal_moves:
                if move.from_square == square_index:
                    moves.append(chess.square_name(move.to_square))
            
            return moves
            
        except ValueError:
            return []
    
    @staticmethod
    def get_piece_attacks(board: chess.Board, square: str) -> List[str]:
        """
        Get all squares attacked by piece at given square
        
        Args:
            board: Chess board instance
            square: Square in algebraic notation
            
        Returns:
            List of attacked squares in algebraic notation
        """
        try:
            square_index = chess.parse_square(square)
            piece = board.piece_at(square_index)
            
            if piece is None:
                return []
            
            attacks = []
            for target_square in chess.SQUARES:
                if board.is_attacked_by(piece.color, target_square):
                    # Check if this specific piece is doing the attacking
                    # This is a simplified check - in a real game you might need more sophisticated logic
                    attacks.append(chess.square_name(target_square))
            
            return attacks
            
        except ValueError:
            return []
    
    @staticmethod
    def get_material_count(board: chess.Board) -> Dict[str, Dict[str, int]]:
        """
        Get material count for both sides
        
        Args:
            board: Chess board instance
            
        Returns:
            Dictionary with material counts for white and black
        """
        white_pieces = {'pawn': 0, 'knight': 0, 'bishop': 0, 'rook': 0, 'queen': 0, 'king': 0}
        black_pieces = {'pawn': 0, 'knight': 0, 'bishop': 0, 'rook': 0, 'queen': 0, 'king': 0}
        
        piece_names = {
            chess.PAWN: 'pawn',
            chess.KNIGHT: 'knight',
            chess.BISHOP: 'bishop',
            chess.ROOK: 'rook',
            chess.QUEEN: 'queen',
            chess.KING: 'king'
        }
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece:
                piece_name = piece_names[piece.piece_type]
                if piece.color == chess.WHITE:
                    white_pieces[piece_name] += 1
                else:
                    black_pieces[piece_name] += 1
        
        return {
            'white': white_pieces,
            'black': black_pieces
        }
    
    @staticmethod
    def get_material_value(board: chess.Board) -> Dict[str, int]:
        """
        Get total material value for both sides
        
        Args:
            board: Chess board instance
            
        Returns:
            Dictionary with material values for white and black
        """
        white_value = 0
        black_value = 0
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece:
                value = ChessPiece.PIECE_VALUES[piece.piece_type]
                if piece.color == chess.WHITE:
                    white_value += value
                else:
                    black_value += value
        
        return {
            'white': white_value,
            'black': black_value,
            'advantage': white_value - black_value
        }
    
    @staticmethod
    def get_piece_list(board: chess.Board) -> Dict[str, List[Dict[str, str]]]:
        """
        Get list of all pieces on the board
        
        Args:
            board: Chess board instance
            
        Returns:
            Dictionary with lists of white and black pieces
        """
        white_pieces = []
        black_pieces = []
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece:
                piece_info = {
                    'type': chess.piece_name(piece.piece_type),
                    'square': chess.square_name(square),
                    'symbol': piece.symbol(),
                    'unicode': ChessPiece.UNICODE_PIECES[piece.color][piece.piece_type]
                }
                
                if piece.color == chess.WHITE:
                    white_pieces.append(piece_info)
                else:
                    black_pieces.append(piece_info)
        
        return {
            'white': white_pieces,
            'black': black_pieces
        }
    
    @staticmethod
    def is_promotion_move(board: chess.Board, move_str: str) -> bool:
        """
        Check if a move is a pawn promotion
        
        Args:
            board: Chess board instance
            move_str: Move in UCI notation
            
        Returns:
            True if the move is a promotion
        """
        try:
            move = chess.Move.from_uci(move_str)
            return move.promotion is not None
        except ValueError:
            return False
    
    @staticmethod
    def get_promotion_pieces() -> List[Dict[str, str]]:
        """
        Get available promotion pieces
        
        Returns:
            List of promotion piece options
        """
        return [
            {'type': 'queen', 'symbol': 'Q', 'name': 'Queen'},
            {'type': 'rook', 'symbol': 'R', 'name': 'Rook'},
            {'type': 'bishop', 'symbol': 'B', 'name': 'Bishop'},
            {'type': 'knight', 'symbol': 'N', 'name': 'Knight'}
        ]