chess_engine/ai/stockfish_wrapper.py

# chess_engine/ai/stockfish_wrapper.py

import chess
import chess.engine
from stockfish import Stockfish
from typing import Optional, Dict, List, Tuple, Any
from enum import Enum
import asyncio
import logging
from dataclasses import dataclass

class DifficultyLevel(Enum):
    BEGINNER = 1
    EASY = 3
    MEDIUM = 5
    HARD = 8
    EXPERT = 12
    MASTER = 15

@dataclass
class EngineConfig:
    """Configuration for Stockfish engine"""
    depth: int = 10
    time_limit: float = 1.0  # seconds
    threads: int = 1
    hash_size: int = 16  # MB
    skill_level: int = 20  # 0-20 (20 is strongest)
    contempt: int = 0
    ponder: bool = False

@dataclass
class MoveAnalysis:
    """Analysis result for a move"""
    best_move: str
    evaluation: float
    depth: int
    principal_variation: List[str]
    time_taken: float
    nodes_searched: int

class StockfishWrapper:
    """
    Wrapper for Stockfish chess engine with both python-chess and stockfish library support
    """
    
    def __init__(self, stockfish_path: Optional[str] = None, config: Optional[EngineConfig] = None):
        """
        Initialize Stockfish wrapper
        
        Args:
            stockfish_path: Path to Stockfish executable
            config: Engine configuration
        """
        self.config = config or EngineConfig()
        self.stockfish_path = stockfish_path
        self.engine = None
        self.stockfish = None
        self.is_initialized = False
        
        # Setup logging
        self.logger = logging.getLogger(__name__)
        
    def initialize(self) -> bool:
        """
        Initialize the Stockfish engine
        
        Returns:
            True if successful, False otherwise
        """
        try:
            # Try to initialize with stockfish library first
            if self.stockfish_path:
                self.stockfish = Stockfish(
                    path=self.stockfish_path,
                    depth=self.config.depth,
                    parameters={
                        "Threads": self.config.threads,
                        "Hash": self.config.hash_size,
                        "Skill Level": self.config.skill_level,
                        "Contempt": self.config.contempt,
                        "Ponder": self.config.ponder
                    }
                )
            else:
                # Use default system Stockfish
                self.stockfish = Stockfish(
                    depth=self.config.depth,
                    parameters={
                        "Threads": self.config.threads,
                        "Hash": self.config.hash_size,
                        "Skill Level": self.config.skill_level,
                        "Contempt": self.config.contempt,
                        "Ponder": self.config.ponder
                    }
                )
            
            # Test if engine is working
            if self.stockfish.is_fen_valid("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"):
                self.is_initialized = True
                self.logger.info("Stockfish engine initialized successfully")
                return True
            else:
                self.logger.error("Stockfish engine test failed")
                return False
                
        except Exception as e:
            self.logger.error(f"Failed to initialize Stockfish: {e}")
            return False
    
    async def initialize_async(self) -> bool:
        """
        Initialize the chess.engine for async operations
        
        Returns:
            True if successful, False otherwise
        """
        try:
            if self.stockfish_path:
                self.engine = await chess.engine.SimpleEngine.popen_uci(self.stockfish_path)
            else:
                # Try common Stockfish paths
                paths = [
                    "/usr/bin/stockfish",
                    "/usr/local/bin/stockfish",
                    "stockfish",
                    "stockfish.exe"
                ]
                
                for path in paths:
                    try:
                        self.engine = await chess.engine.SimpleEngine.popen_uci(path)
                        break
                    except FileNotFoundError:
                        continue
                
                if not self.engine:
                    raise FileNotFoundError("Stockfish executable not found")
            
            # Configure engine
            await self.engine.configure({
                "Threads": self.config.threads,
                "Hash": self.config.hash_size,
                "Skill Level": self.config.skill_level,
                "Contempt": self.config.contempt,
                "Ponder": self.config.ponder
            })
            
            self.is_initialized = True
            self.logger.info("Async Stockfish engine initialized successfully")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to initialize async Stockfish: {e}")
            return False
    
    def get_best_move(self, board: chess.Board, time_limit: Optional[float] = None) -> Optional[str]:
        """
        Get the best move for current position
        
        Args:
            board: Current chess board position
            time_limit: Time limit in seconds (overrides config)
            
        Returns:
            Best move in UCI notation or None if error
        """
        if not self.is_initialized or not self.stockfish:
            return None
        
        try:
            # Set position
            self.stockfish.set_fen_position(board.fen())
            
            # Get best move
            best_move = self.stockfish.get_best_move_time(
                time_limit or self.config.time_limit * 1000  # Convert to milliseconds
            )
            
            return best_move
            
        except Exception as e:
            self.logger.error(f"Error getting best move: {e}")
            return None
    
    async def get_best_move_async(self, board: chess.Board, time_limit: Optional[float] = None) -> Optional[MoveAnalysis]:
        """
        Get the best move asynchronously with detailed analysis
        
        Args:
            board: Current chess board position
            time_limit: Time limit in seconds
            
        Returns:
            MoveAnalysis object or None if error
        """
        if not self.is_initialized or not self.engine:
            return None
        
        try:
            # Set up analysis parameters
            limit = chess.engine.Limit(
                time=time_limit or self.config.time_limit,
                depth=self.config.depth
            )
            
            # Analyze position
            info = await self.engine.analyse(board, limit)
            
            # Get best move
            result = await self.engine.play(board, limit)
            
            # Extract information
            evaluation = info.get("score", chess.engine.Cp(0))
            pv = info.get("pv", [])
            
            # Convert evaluation to numerical value
            if evaluation.is_mate():
                eval_value = 1000.0 if evaluation.mate() > 0 else -1000.0
            else:
                eval_value = evaluation.score() / 100.0  # Convert centipawns to pawns
            
            return MoveAnalysis(
                best_move=result.move.uci() if result.move else "",
                evaluation=eval_value,
                depth=info.get("depth", 0),
                principal_variation=[move.uci() for move in pv],
                time_taken=info.get("time", 0.0),
                nodes_searched=info.get("nodes", 0)
            )
            
        except Exception as e:
            self.logger.error(f"Error in async move analysis: {e}")
            return None
    
    def evaluate_position(self, board: chess.Board) -> Optional[float]:
        """
        Evaluate current position
        
        Args:
            board: Current chess board position
            
        Returns:
            Evaluation in pawns (positive for white advantage)
        """
        if not self.is_initialized or not self.stockfish:
            return None
        
        try:
            self.stockfish.set_fen_position(board.fen())
            evaluation = self.stockfish.get_evaluation()
            
            if evaluation is None:
                return 0.0
            
            if evaluation["type"] == "cp":
                return evaluation["value"] / 100.0  # Convert centipawns to pawns
            elif evaluation["type"] == "mate":
                return 1000.0 if evaluation["value"] > 0 else -1000.0
            
            return 0.0
            
        except Exception as e:
            self.logger.error(f"Error evaluating position: {e}")
            return None
    
    def get_legal_moves_with_evaluation(self, board: chess.Board) -> List[Tuple[str, float]]:
        """
        Get all legal moves with their evaluations
        
        Args:
            board: Current chess board position
            
        Returns:
            List of (move, evaluation) tuples
        """
        if not self.is_initialized or not self.stockfish:
            return []
        
        moves_with_eval = []
        
        try:
            for move in board.legal_moves:
                # Make move temporarily
                board.push(move)
                
                # Evaluate position
                evaluation = self.evaluate_position(board)
                
                # Undo move
                board.pop()
                
                if evaluation is not None:
                    moves_with_eval.append((move.uci(), -evaluation))  # Negate for opponent's perspective
            
            # Sort by evaluation (best first)
            moves_with_eval.sort(key=lambda x: x[1], reverse=True)
            
            return moves_with_eval
            
        except Exception as e:
            self.logger.error(f"Error getting moves with evaluation: {e}")
            return []
    
    def set_difficulty_level(self, level: DifficultyLevel):
        """
        Set AI difficulty level
        
        Args:
            level: Difficulty level enum
        """
        skill_levels = {
            DifficultyLevel.BEGINNER: 1,
            DifficultyLevel.EASY: 3,
            DifficultyLevel.MEDIUM: 8,
            DifficultyLevel.HARD: 12,
            DifficultyLevel.EXPERT: 17,
            DifficultyLevel.MASTER: 20
        }
        
        depths = {
            DifficultyLevel.BEGINNER: 3,
            DifficultyLevel.EASY: 5,
            DifficultyLevel.MEDIUM: 8,
            DifficultyLevel.HARD: 12,
            DifficultyLevel.EXPERT: 15,
            DifficultyLevel.MASTER: 20
        }
        
        self.config.skill_level = skill_levels[level]
        self.config.depth = depths[level]
        
        # Update engine parameters if initialized
        if self.is_initialized and self.stockfish:
            self.stockfish.set_depth(self.config.depth)
            self.stockfish.set_skill_level(self.config.skill_level)
    
    def get_engine_info(self) -> Dict[str, Any]:
        """
        Get engine information and statistics
        
        Returns:
            Dictionary with engine info
        """
        if not self.is_initialized:
            return {"status": "not_initialized"}
        
        return {
            "status": "initialized",
            "config": {
                "depth": self.config.depth,
                "skill_level": self.config.skill_level,
                "time_limit": self.config.time_limit,
                "threads": self.config.threads,
                "hash_size": self.config.hash_size
            },
            "stockfish_available": self.stockfish is not None,
            "async_engine_available": self.engine is not None
        }
    
    def is_move_blunder(self, board: chess.Board, move: str, threshold: float = 2.0) -> bool:
        """
        Check if a move is a blunder
        
        Args:
            board: Current chess board position
            move: Move to check in UCI notation
            threshold: Evaluation drop threshold for blunder detection
            
        Returns:
            True if move is a blunder
        """
        if not self.is_initialized:
            return False
        
        try:
            # Get current position evaluation
            current_eval = self.evaluate_position(board)
            if current_eval is None:
                return False
            
            # Make the move
            move_obj = chess.Move.from_uci(move)
            if move_obj not in board.legal_moves:
                return True  # Illegal move is definitely a blunder
            
            board.push(move_obj)
            
            # Get evaluation after move
            new_eval = self.evaluate_position(board)
            
            # Undo move
            board.pop()
            
            if new_eval is None:
                return False
            
            # Check if evaluation dropped significantly
            # Note: negate new_eval because it's opponent's turn
            eval_drop = current_eval - (-new_eval)
            
            return eval_drop > threshold
            
        except Exception as e:
            self.logger.error(f"Error checking blunder: {e}")
            return False
    
    def close(self):
        """Close the engine connection"""
        if self.engine:
            asyncio.create_task(self.engine.quit())
        
        self.is_initialized = False
        self.logger.info("Stockfish engine closed")
    
    def __enter__(self):
        """Context manager entry"""
        self.initialize()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        self.close()