chess_engine/ai/evaluation.py

# chess_engine/ai/evaluation.py

import chess
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
import math

@dataclass
class PositionEvaluation:
    """Complete position evaluation data"""
    total_score: float
    material_score: float
    positional_score: float
    safety_score: float
    mobility_score: float
    pawn_structure_score: float
    endgame_score: float
    white_advantage: float
    evaluation_breakdown: Dict[str, float]

class ChessEvaluator:
    """
    Chess position evaluator with various strategic factors
    """
    
    # Piece values in centipawns
    PIECE_VALUES = {
        chess.PAWN: 100,
        chess.KNIGHT: 320,
        chess.BISHOP: 330,
        chess.ROOK: 500,
        chess.QUEEN: 900,
        chess.KING: 20000
    }
    
    # Piece-square tables for positional evaluation
    PAWN_TABLE = [
        [0,  0,  0,  0,  0,  0,  0,  0],
        [50, 50, 50, 50, 50, 50, 50, 50],
        [10, 10, 20, 30, 30, 20, 10, 10],
        [5,  5, 10, 25, 25, 10,  5,  5],
        [0,  0,  0, 20, 20,  0,  0,  0],
        [5, -5,-10,  0,  0,-10, -5,  5],
        [5, 10, 10,-20,-20, 10, 10,  5],
        [0,  0,  0,  0,  0,  0,  0,  0]
    ]
    
    KNIGHT_TABLE = [
        [-50,-40,-30,-30,-30,-30,-40,-50],
        [-40,-20,  0,  0,  0,  0,-20,-40],
        [-30,  0, 10, 15, 15, 10,  0,-30],
        [-30,  5, 15, 20, 20, 15,  5,-30],
        [-30,  0, 15, 20, 20, 15,  0,-30],
        [-30,  5, 10, 15, 15, 10,  5,-30],
        [-40,-20,  0,  5,  5,  0,-20,-40],
        [-50,-40,-30,-30,-30,-30,-40,-50]
    ]
    
    BISHOP_TABLE = [
        [-20,-10,-10,-10,-10,-10,-10,-20],
        [-10,  0,  0,  0,  0,  0,  0,-10],
        [-10,  0,  5, 10, 10,  5,  0,-10],
        [-10,  5,  5, 10, 10,  5,  5,-10],
        [-10,  0, 10, 10, 10, 10,  0,-10],
        [-10, 10, 10, 10, 10, 10, 10,-10],
        [-10,  5,  0,  0,  0,  0,  5,-10],
        [-20,-10,-10,-10,-10,-10,-10,-20]
    ]
    
    ROOK_TABLE = [
        [0,  0,  0,  0,  0,  0,  0,  0],
        [5, 10, 10, 10, 10, 10, 10,  5],
        [-5,  0,  0,  0,  0,  0,  0, -5],
        [-5,  0,  0,  0,  0,  0,  0, -5],
        [-5,  0,  0,  0,  0,  0,  0, -5],
        [-5,  0,  0,  0,  0,  0,  0, -5],
        [-5,  0,  0,  0,  0,  0,  0, -5],
        [0,  0,  0,  5,  5,  0,  0,  0]
    ]
    
    QUEEN_TABLE = [
        [-20,-10,-10, -5, -5,-10,-10,-20],
        [-10,  0,  0,  0,  0,  0,  0,-10],
        [-10,  0,  5,  5,  5,  5,  0,-10],
        [-5,  0,  5,  5,  5,  5,  0, -5],
        [0,  0,  5,  5,  5,  5,  0, -5],
        [-10,  5,  5,  5,  5,  5,  0,-10],
        [-10,  0,  5,  0,  0,  0,  0,-10],
        [-20,-10,-10, -5, -5,-10,-10,-20]
    ]
    
    KING_TABLE_MIDDLEGAME = [
        [-30,-40,-40,-50,-50,-40,-40,-30],
        [-30,-40,-40,-50,-50,-40,-40,-30],
        [-30,-40,-40,-50,-50,-40,-40,-30],
        [-30,-40,-40,-50,-50,-40,-40,-30],
        [-20,-30,-30,-40,-40,-30,-30,-20],
        [-10,-20,-20,-20,-20,-20,-20,-10],
        [20, 20,  0,  0,  0,  0, 20, 20],
        [20, 30, 10,  0,  0, 10, 30, 20]
    ]
    
    KING_TABLE_ENDGAME = [
        [-50,-40,-30,-20,-20,-30,-40,-50],
        [-30,-20,-10,  0,  0,-10,-20,-30],
        [-30,-10, 20, 30, 30, 20,-10,-30],
        [-30,-10, 30, 40, 40, 30,-10,-30],
        [-30,-10, 30, 40, 40, 30,-10,-30],
        [-30,-10, 20, 30, 30, 20,-10,-30],
        [-30,-30,  0,  0,  0,  0,-30,-30],
        [-50,-30,-30,-30,-30,-30,-30,-50]
    ]
    
    def __init__(self):
        """Initialize the evaluator"""
        self.piece_square_tables = {
            chess.PAWN: self.PAWN_TABLE,
            chess.KNIGHT: self.KNIGHT_TABLE,
            chess.BISHOP: self.BISHOP_TABLE,
            chess.ROOK: self.ROOK_TABLE,
            chess.QUEEN: self.QUEEN_TABLE,
            chess.KING: self.KING_TABLE_MIDDLEGAME
        }
    
    def evaluate_position(self, board: chess.Board) -> PositionEvaluation:
        """
        Comprehensive position evaluation
        
        Args:
            board: Chess board to evaluate
            
        Returns:
            PositionEvaluation object with detailed analysis
        """
        if board.is_checkmate():
            score = -20000 if board.turn == chess.WHITE else 20000
            return PositionEvaluation(
                total_score=score,
                material_score=score,
                positional_score=0,
                safety_score=0,
                mobility_score=0,
                pawn_structure_score=0,
                endgame_score=0,
                white_advantage=score,
                evaluation_breakdown={"checkmate": score}
            )
        
        if board.is_stalemate() or board.is_insufficient_material():
            return PositionEvaluation(
                total_score=0,
                material_score=0,
                positional_score=0,
                safety_score=0,
                mobility_score=0,
                pawn_structure_score=0,
                endgame_score=0,
                white_advantage=0,
                evaluation_breakdown={"draw": 0}
            )
        
        # Calculate individual evaluation components
        material_score = self._evaluate_material(board)
        positional_score = self._evaluate_position_tables(board)
        safety_score = self._evaluate_king_safety(board)
        mobility_score = self._evaluate_mobility(board)
        pawn_structure_score = self._evaluate_pawn_structure(board)
        endgame_score = self._evaluate_endgame_factors(board)
        
        # Combine scores
        total_score = (
            material_score + 
            positional_score + 
            safety_score + 
            mobility_score + 
            pawn_structure_score + 
            endgame_score
        )
        
        # Create breakdown
        breakdown = {
            "material": material_score,
            "positional": positional_score,
            "safety": safety_score,
            "mobility": mobility_score,
            "pawn_structure": pawn_structure_score,
            "endgame": endgame_score
        }
        
        return PositionEvaluation(
            total_score=total_score,
            material_score=material_score,
            positional_score=positional_score,
            safety_score=safety_score,
            mobility_score=mobility_score,
            pawn_structure_score=pawn_structure_score,
            endgame_score=endgame_score,
            white_advantage=total_score,
            evaluation_breakdown=breakdown
        )
    
    def _evaluate_material(self, board: chess.Board) -> float:
        """Evaluate material balance"""
        score = 0
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece:
                value = self.PIECE_VALUES[piece.piece_type]
                score += value if piece.color == chess.WHITE else -value
        
        return score
    
    def _evaluate_position_tables(self, board: chess.Board) -> float:
        """Evaluate piece positions using piece-square tables"""
        score = 0
        is_endgame = self._is_endgame(board)
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece:
                rank = chess.square_rank(square)
                file = chess.square_file(square)
                
                # Choose appropriate table for king
                if piece.piece_type == chess.KING:
                    table = self.KING_TABLE_ENDGAME if is_endgame else self.KING_TABLE_MIDDLEGAME
                else:
                    table = self.piece_square_tables[piece.piece_type]
                
                # Flip table for black pieces
                if piece.color == chess.WHITE:
                    value = table[rank][file]
                else:
                    value = -table[7-rank][file]
                
                score += value
        
        return score
    
    def _evaluate_king_safety(self, board: chess.Board) -> float:
        """Evaluate king safety"""
        score = 0
        
        # Check for king exposure
        for color in [chess.WHITE, chess.BLACK]:
            king_square = board.king(color)
            if king_square is None:
                continue
            
            # Count attackers around king
            attackers = 0
            defenders = 0
            
            for square in chess.SQUARES:
                if chess.square_distance(king_square, square) <= 2:
                    if board.is_attacked_by(not color, square):
                        attackers += 1
                    if board.is_attacked_by(color, square):
                        defenders += 1
            
            safety = (defenders - attackers) * 10
            score += safety if color == chess.WHITE else -safety
        
        return score
    
    def _evaluate_mobility(self, board: chess.Board) -> float:
        """Evaluate piece mobility"""
        white_mobility = len(list(board.legal_moves)) if board.turn == chess.WHITE else 0
        
        # Switch turn to count black mobility
        board.turn = not board.turn
        black_mobility = len(list(board.legal_moves)) if board.turn == chess.BLACK else 0
        board.turn = not board.turn  # Switch back
        
        return (white_mobility - black_mobility) * 1.5
    
    def _evaluate_pawn_structure(self, board: chess.Board) -> float:
        """Evaluate pawn structure"""
        score = 0
        
        # Get pawn positions
        white_pawns = [sq for sq in chess.SQUARES if board.piece_at(sq) and 
                      board.piece_at(sq).piece_type == chess.PAWN and 
                      board.piece_at(sq).color == chess.WHITE]
        
        black_pawns = [sq for sq in chess.SQUARES if board.piece_at(sq) and 
                      board.piece_at(sq).piece_type == chess.PAWN and 
                      board.piece_at(sq).color == chess.BLACK]
        
        # Evaluate doubled pawns
        score += self._evaluate_doubled_pawns(white_pawns, chess.WHITE)
        score += self._evaluate_doubled_pawns(black_pawns, chess.BLACK)
        
        # Evaluate isolated pawns
        score += self._evaluate_isolated_pawns(white_pawns, chess.WHITE)
        score += self._evaluate_isolated_pawns(black_pawns, chess.BLACK)
        
        # Evaluate passed pawns
        score += self._evaluate_passed_pawns(board, white_pawns, chess.WHITE)
        score += self._evaluate_passed_pawns(board, black_pawns, chess.BLACK)
        
        return score
    
    def _evaluate_doubled_pawns(self, pawns: List[int], color: chess.Color) -> float:
        """Evaluate doubled pawns penalty"""
        files = {}
        for pawn in pawns:
            file = chess.square_file(pawn)
            files[file] = files.get(file, 0) + 1
        
        doubled_count = sum(max(0, count - 1) for count in files.values())
        penalty = doubled_count * -20
        
        return penalty if color == chess.WHITE else -penalty
    
    def _evaluate_isolated_pawns(self, pawns: List[int], color: chess.Color) -> float:
        """Evaluate isolated pawns penalty"""
        files = set(chess.square_file(pawn) for pawn in pawns)
        isolated_count = 0
        
        for file in files:
            if (file - 1 not in files) and (file + 1 not in files):
                isolated_count += 1
        
        penalty = isolated_count * -15
        return penalty if color == chess.WHITE else -penalty
    
    def _evaluate_passed_pawns(self, board: chess.Board, pawns: List[int], color: chess.Color) -> float:
        """Evaluate passed pawns bonus"""
        bonus = 0
        opponent_color = not color
        
        for pawn in pawns:
            file = chess.square_file(pawn)
            rank = chess.square_rank(pawn)
            
            # Check if pawn is passed
            is_passed = True
            direction = 1 if color == chess.WHITE else -1
            
            # Check files that could block this pawn
            for check_file in [file - 1, file, file + 1]:
                if 0 <= check_file <= 7:
                    for check_rank in range(rank + direction, 8 if color == chess.WHITE else -1, direction):
                        if 0 <= check_rank <= 7:
                            square = chess.square(check_file, check_rank)
                            piece = board.piece_at(square)
                            if piece and piece.piece_type == chess.PAWN and piece.color == opponent_color:
                                is_passed = False
                                break
                    if not is_passed:
                        break
            
            if is_passed:
                # Bonus increases with advancement
                advancement = rank if color == chess.WHITE else 7 - rank
                bonus += advancement * 10
        
        return bonus if color == chess.WHITE else -bonus
    
    def _is_endgame(self, board: chess.Board) -> bool:
        """
        Determine if the position is in the endgame
        
        Args:
            board: Chess board to evaluate
            
        Returns:
            True if position is in endgame, False otherwise
        """
        # Count major pieces (queens and rooks)
        queens = 0
        rooks = 0
        total_material = 0
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece:
                if piece.piece_type == chess.QUEEN:
                    queens += 1
                elif piece.piece_type == chess.ROOK:
                    rooks += 1
                total_material += self.PIECE_VALUES[piece.piece_type]
        
        # Endgame conditions:
        # 1. No queens
        # 2. Only one queen total and no other major pieces
        # 3. Less than 25% of starting material
        return (queens == 0) or (queens == 1 and rooks <= 1) or (total_material < 3200)
    
    def _evaluate_endgame_factors(self, board: chess.Board) -> float:
        """Evaluate endgame-specific factors"""
        if not self._is_endgame(board):
            return 0
        
        score = 0
        
        # King centralization in endgame
        score += self._evaluate_king_centralization(board)
        
        # Passed pawns become more valuable in endgame
        score += self._evaluate_endgame_passed_pawns(board)
        
        # Rook on open files
        score += self._evaluate_rooks_on_open_files(board)
        
        return score
    
    def _evaluate_king_centralization(self, board: chess.Board) -> float:
        """
        Evaluate king centralization in endgame
        Kings should move to the center in endgames
        """
        score = 0
        
        for color in [chess.WHITE, chess.BLACK]:
            king_square = board.king(color)
            if king_square is None:
                continue
            
            # Calculate distance from center (d4, d5, e4, e5)
            file = chess.square_file(king_square)
            rank = chess.square_rank(king_square)
            
            # Distance from center files (0-3.5)
            file_distance = abs(3.5 - file)
            
            # Distance from center ranks (0-3.5)
            rank_distance = abs(3.5 - rank)
            
            # Manhattan distance from center
            center_distance = file_distance + rank_distance
            
            # Bonus for being close to center (max 15 points)
            centralization_bonus = (7 - center_distance) * 3
            
            if color == chess.WHITE:
                score += centralization_bonus
            else:
                score -= centralization_bonus
        
        return score
    
    def _evaluate_endgame_passed_pawns(self, board: chess.Board) -> float:
        """
        Evaluate passed pawns in endgame - they're more valuable
        """
        score = 0
        
        # Get pawn positions
        white_pawns = [sq for sq in chess.SQUARES if board.piece_at(sq) and 
                      board.piece_at(sq).piece_type == chess.PAWN and 
                      board.piece_at(sq).color == chess.WHITE]
        
        black_pawns = [sq for sq in chess.SQUARES if board.piece_at(sq) and 
                      board.piece_at(sq).piece_type == chess.PAWN and 
                      board.piece_at(sq).color == chess.BLACK]
        
        # Check for passed pawns
        for pawn in white_pawns:
            if self._is_passed_pawn(board, pawn, chess.WHITE):
                rank = chess.square_rank(pawn)
                # Bonus increases dramatically with advancement
                bonus = (rank * rank) * 5
                score += bonus
        
        for pawn in black_pawns:
            if self._is_passed_pawn(board, pawn, chess.BLACK):
                rank = 7 - chess.square_rank(pawn)  # Flip for black
                # Bonus increases dramatically with advancement
                bonus = (rank * rank) * 5
                score -= bonus
        
        return score
    
    def _is_passed_pawn(self, board: chess.Board, square: int, color: chess.Color) -> bool:
        """Check if a pawn is passed"""
        file = chess.square_file(square)
        rank = chess.square_rank(square)
        
        # Direction of pawn movement
        direction = 1 if color == chess.WHITE else -1
        
        # Check files that could block this pawn
        for check_file in [file - 1, file, file + 1]:
            if 0 <= check_file <= 7:
                for check_rank in range(rank + direction, 8 if color == chess.WHITE else -1, direction):
                    if 0 <= check_rank <= 7:
                        check_square = chess.square(check_file, check_rank)
                        piece = board.piece_at(check_square)
                        if piece and piece.piece_type == chess.PAWN and piece.color != color:
                            return False
        
        return True
    
    def _evaluate_rooks_on_open_files(self, board: chess.Board) -> float:
        """Evaluate rooks on open or semi-open files"""
        score = 0
        
        # Get all files with pawns
        files_with_white_pawns = set()
        files_with_black_pawns = set()
        
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece and piece.piece_type == chess.PAWN:
                file = chess.square_file(square)
                if piece.color == chess.WHITE:
                    files_with_white_pawns.add(file)
                else:
                    files_with_black_pawns.add(file)
        
        # Check rooks
        for square in chess.SQUARES:
            piece = board.piece_at(square)
            if piece and piece.piece_type == chess.ROOK:
                file = chess.square_file(square)
                
                # Open file (no pawns)
                if file not in files_with_white_pawns and file not in files_with_black_pawns:
                    bonus = 25
                # Semi-open file (no friendly pawns)
                elif (piece.color == chess.WHITE and file not in files_with_white_pawns) or \
                     (piece.color == chess.BLACK and file not in files_with_black_pawns):
                    bonus = 15
                else:
                    bonus = 0
                
                if piece.color == chess.WHITE:
                    score += bonus
                else:
                    score -= bonus
        
        return score