engine.py

"""
Deep Search Engine for Synapse-Base
Alpha-Beta pruning with advanced move ordering
CPU-optimized for HF Spaces (2 vCPU, 16GB RAM)
"""

import chess
import time
import logging
from typing import Optional, Tuple, List
from model_loader import SynapseModel

logger = logging.getLogger(__name__)


class SynapseEngine:
    """
    Chess engine with neural network evaluation and alpha-beta search
    """
    
    def __init__(self, model_path: str, num_threads: int = 2):
        """
        Initialize engine
        
        Args:
            model_path: Path to ONNX model
            num_threads: CPU threads for inference
        """
        self.model = SynapseModel(model_path, num_threads)
        self.nodes_evaluated = 0
        self.cache = {}  # Transposition table
        self.max_cache_size = 50000  # Limit cache to ~400MB RAM
        
        # Move ordering weights
        self.MVV_LVA = {
            (chess.PAWN, chess.QUEEN): 50,
            (chess.PAWN, chess.ROOK): 40,
            (chess.KNIGHT, chess.QUEEN): 45,
            (chess.KNIGHT, chess.ROOK): 35,
            (chess.BISHOP, chess.QUEEN): 45,
            (chess.BISHOP, chess.ROOK): 35,
            (chess.ROOK, chess.QUEEN): 40,
            (chess.QUEEN, chess.PAWN): -10,
        }
        
        logger.info("🎯 Engine initialized")
    
    def evaluate_position(self, board: chess.Board) -> float:
        """
        Evaluate position using neural network
        
        Args:
            board: chess.Board object
        
        Returns:
            float: evaluation score from white's perspective
        """
        self.nodes_evaluated += 1
        
        # Check cache
        fen_key = board.fen().split(' ')[0]  # Position only (no move counters)
        if fen_key in self.cache:
            return self.cache[fen_key]
        
        # Get evaluation from model
        result = self.model.evaluate(board.fen())
        evaluation = result['value']
        
        # Flip score for black's turn
        if board.turn == chess.BLACK:
            evaluation = -evaluation
        
        # Cache result
        if len(self.cache) < self.max_cache_size:
            self.cache[fen_key] = evaluation
        
        return evaluation
    
    def order_moves(self, board: chess.Board, moves: List[chess.Move]) -> List[chess.Move]:
        """
        Order moves for better alpha-beta pruning
        
        Priority:
        1. Captures (MVV-LVA)
        2. Checks
        3. Promotions
        4. Castling
        5. Other moves
        
        Args:
            board: Current position
            moves: List of legal moves
        
        Returns:
            Sorted list of moves
        """
        scored_moves = []
        
        for move in moves:
            score = 0
            
            # === HIGH PRIORITY ===
            
            # Captures (Most Valuable Victim - Least Valuable Attacker)
            if board.is_capture(move):
                captured_piece = board.piece_at(move.to_square)
                moving_piece = board.piece_at(move.from_square)
                
                if captured_piece and moving_piece:
                    # Base capture value
                    piece_values = {
                        chess.PAWN: 1,
                        chess.KNIGHT: 3,
                        chess.BISHOP: 3,
                        chess.ROOK: 5,
                        chess.QUEEN: 9,
                        chess.KING: 0
                    }
                    
                    victim_value = piece_values.get(captured_piece.piece_type, 0)
                    attacker_value = piece_values.get(moving_piece.piece_type, 1)
                    
                    # MVV-LVA: prioritize high-value captures with low-value pieces
                    score += (victim_value * 10 - attacker_value) * 1000
            
            # Promotions
            if move.promotion:
                if move.promotion == chess.QUEEN:
                    score += 9000
                elif move.promotion == chess.KNIGHT:
                    score += 3000
            
            # Checks
            board.push(move)
            if board.is_check():
                score += 5000
            board.pop()
            
            # Castling
            if board.is_castling(move):
                score += 2000
            
            # === MEDIUM PRIORITY ===
            
            # Center control
            center_squares = [chess.D4, chess.D5, chess.E4, chess.E5]
            if move.to_square in center_squares:
                score += 100
            
            # Piece development (knights and bishops to good squares)
            piece = board.piece_at(move.from_square)
            if piece and piece.piece_type in [chess.KNIGHT, chess.BISHOP]:
                # Bonus for moving from back rank
                if move.from_square // 8 in [0, 7]:
                    score += 50
            
            scored_moves.append((score, move))
        
        # Sort by score (descending)
        scored_moves.sort(key=lambda x: x[0], reverse=True)
        
        return [move for score, move in scored_moves]
    
    def quiescence_search(
        self,
        board: chess.Board,
        alpha: float,
        beta: float,
        depth: int = 6
    ) -> float:
        """
        Quiescence search to avoid horizon effect
        Only searches captures and checks
        
        Args:
            board: Current position
            alpha: Alpha value
            beta: Beta value
            depth: Remaining quiescence depth
        
        Returns:
            Evaluation score
        """
        # Stand-pat evaluation
        stand_pat = self.evaluate_position(board)
        
        if stand_pat >= beta:
            return beta
        if alpha < stand_pat:
            alpha = stand_pat
        
        if depth == 0:
            return stand_pat
        
        # Only search tactical moves (captures and checks)
        tactical_moves = [
            move for move in board.legal_moves
            if board.is_capture(move) or board.gives_check(move)
        ]
        
        if not tactical_moves:
            return stand_pat
        
        tactical_moves = self.order_moves(board, tactical_moves)
        
        for move in tactical_moves:
            board.push(move)
            score = -self.quiescence_search(board, -beta, -alpha, depth - 1)
            board.pop()
            
            if score >= beta:
                return beta
            if score > alpha:
                alpha = score
        
        return alpha
    
    def alpha_beta(
        self,
        board: chess.Board,
        depth: int,
        alpha: float,
        beta: float,
        start_time: float,
        time_limit: float
    ) -> Tuple[float, Optional[chess.Move]]:
        """
        Alpha-beta pruning search
        
        Args:
            board: Current position
            depth: Remaining search depth
            alpha: Alpha value
            beta: Beta value
            start_time: Search start time
            time_limit: Max time in seconds
        
        Returns:
            (evaluation, best_move)
        """
        # Time check
        if time.time() - start_time > time_limit:
            return self.evaluate_position(board), None
        
        # Terminal node checks
        if board.is_game_over():
            if board.is_checkmate():
                return -10000, None  # Mate is worst
            return 0, None  # Draw
        
        # Depth limit reached - use quiescence search
        if depth == 0:
            return self.quiescence_search(board, alpha, beta, depth=2), None
        
        legal_moves = list(board.legal_moves)
        if not legal_moves:
            return 0, None
        
        # Move ordering
        ordered_moves = self.order_moves(board, legal_moves)
        
        best_move = ordered_moves[0]
        best_score = float('-inf')
        
        for move in ordered_moves:
            board.push(move)
            
            # Recursive search
            score, _ = self.alpha_beta(
                board, depth - 1, -beta, -alpha, start_time, time_limit
            )
            score = -score
            
            board.pop()
            
            # Update best move
            if score > best_score:
                best_score = score
                best_move = move
            
            # Alpha-beta pruning
            alpha = max(alpha, score)
            if alpha >= beta:
                break  # Beta cutoff
        
        return best_score, best_move
    
    def get_best_move(
        self,
        fen: str,
        depth: int = 6,
        time_limit: int = 5000
    ) -> dict:
        """
        Get best move for position
        
        Args:
            fen: FEN string
            depth: Search depth
            time_limit: Time limit in milliseconds
        
        Returns:
            dict with best_move, evaluation, nodes, etc.
        """
        board = chess.Board(fen)
        
        # Reset counters
        self.nodes_evaluated = 0
        
        # Time limit in seconds
        time_limit_sec = time_limit / 1000.0
        start_time = time.time()
        
        # Special case: only one legal move
        legal_moves = list(board.legal_moves)
        if len(legal_moves) == 1:
            return {
                'best_move': legal_moves[0].uci(),
                'evaluation': self.evaluate_position(board),
                'depth_searched': 0,
                'nodes_evaluated': 1,
                'pv': [legal_moves[0].uci()]
            }
        
        # Iterative deepening
        best_move = None
        best_eval = None
        
        for current_depth in range(1, depth + 1):
            # Check time
            if time.time() - start_time > time_limit_sec * 0.9:
                break
            
            try:
                eval_score, move = self.alpha_beta(
                    board,
                    current_depth,
                    float('-inf'),
                    float('inf'),
                    start_time,
                    time_limit_sec
                )
                
                if move:
                    best_move = move
                    best_eval = eval_score
                
            except Exception as e:
                logger.warning(f"Search error at depth {current_depth}: {e}")
                break
        
        # Fallback to first legal move if search failed
        if best_move is None:
            best_move = legal_moves[0]
            best_eval = self.evaluate_position(board)
        
        return {
            'best_move': best_move.uci(),
            'evaluation': round(best_eval, 4),
            'depth_searched': current_depth,
            'nodes_evaluated': self.nodes_evaluated,
            'pv': [best_move.uci()]
        }
    
    def validate_fen(self, fen: str) -> bool:
        """Validate FEN string"""
        try:
            chess.Board(fen)
            return True
        except:
            return False
    
    def get_model_size(self) -> float:
        """Get model size in MB"""
        return self.model.get_size_mb()