stockfish.py

import math
import serial
import chess
import chess.engine

# Define the serial port for Arduino connection
def connect_to_robot(port, baudrate=9600):
    try:
        ser = serial.Serial(port, baudrate, timeout=1)
        print(f"Connected to robot on {port}")
        return ser
    except Exception as e:
        print(f"Failed to connect to {port}: {e}")
        return None

SQUARE_SIZE = 3.2
offset = 12

filemap = {
    "h": 3.5 + 0.2, "g": 2.5 + 0.1, "f": 1.5, "e": 0.5 - 0.2, "d": -0.5 - 0.3, "c": -1.5, "b": -2.5, "a": -3.5
}

class ChessRobotArm:
    def __init__(self, a: int, b: int) -> None:
        self.a = a  # Length of arm a (cm)
        self.b = b  # Length of arm b (cm)

    def base_angle(self, x: float, y: float):
        y += offset / SQUARE_SIZE + 0.3
        d = math.sqrt((x * SQUARE_SIZE)**2 + (y * SQUARE_SIZE)**2)
        angle = math.atan(y / x) * 180 / math.pi
        if angle < 0:
            angle += 180
        if x < 0:
            angle -= x * 3
        return d, angle

    def inverse_kinematics(self, x: float, y: float):
        try:
            elbow_theta = math.acos((x**2 + y**2 - self.a**2 - self.b**2) / (2 * self.a * self.b))
            shoulder_theta = math.atan(y / x) - math.atan((self.b * math.sin(elbow_theta)) / (self.a + self.b * math.cos(elbow_theta)))
            return -shoulder_theta * 180 / math.pi, 180 - elbow_theta * 180 / math.pi
        except:
            pass

    def go_to(self, square: str):
        x = filemap[square[0]]
        y = int(square[1]) + 0.5
        b = self.base_angle(x, y)
        d = b[0]
        rank = int(square[1])
        vertical_offset = -5 + rank / 2.6
        b_offset = 1
        s, e = self.inverse_kinematics(vertical_offset, d + rank / 10)
        e -= 3
        return b[1] + b_offset, s - 90, e


def get_stockfish_move(board):
    # Specify the correct path for Stockfish engine
    engine_path = "C:/users/ASUS/Desktop/stockfish/stockfish-windows-x86-64-avx2.exe"
    engine = chess.engine.SimpleEngine.popen_uci(engine_path)
    
    # Get Stockfish's best move (White's turn)
    result = engine.play(board, chess.engine.Limit(time=2.0))  # Time limit for move
    move = result.move
    engine.quit()
    return move


if __name__ == '__main__':
    robot = ChessRobotArm(20, 20)
    serial_connection = connect_to_robot("COM9")
    if not serial_connection:
        print("Unable to establish a connection. Exiting.")
        exit()

    board = chess.Board()

    while not board.is_game_over(): 
        print(board)
        # Get the move from Stockfish (e.g., 'e2e4')
        stockfish_move = get_stockfish_move(board)
        print(f"Stockfish suggests move: {stockfish_move}")
        
        # Push Stockfish's move to the board
        board.push(stockfish_move)
        print(f"Board after Stockfish's move: {board}")

        # Extract the from and to squares of the move (e.g., 'e2' and 'e4')
        from_square = str(stockfish_move)[:2]
        to_square = str(stockfish_move)[2:]

        # Move the robot arm to the 'from' square first
        base_angle, shoulder_angle, elbow_angle = robot.go_to(from_square)
        print(f"Moving to {from_square} - Base Angle: {base_angle}, Shoulder Angle: {180 - shoulder_angle}, Elbow Angle: {elbow_angle}")
        command = f"B:{base_angle:.2f} S:{180 - shoulder_angle:.2f} E:{elbow_angle:.2f}\n"
        serial_connection.write(command.encode())

        # Move the robot arm to the 'to' square (final move)
        base_angle, shoulder_angle, elbow_angle = robot.go_to(to_square)
        print(f"Moving to {to_square} - Base Angle: {base_angle}, Shoulder Angle: {180 - shoulder_angle}, Elbow Angle: {elbow_angle}")
        command = f"B:{base_angle:.2f} S:{180 - shoulder_angle:.2f} E:{elbow_angle:.2f}\n"
        serial_connection.write(command.encode())

        # Wait for robot to finish before continuing
        input("Press Enter to continue after move completion.")

        # Prompt for user move (Black's turn)
        user_move = input("Your move (in UCI format, e.g., e7e5): ")

        try:
            # Validate and push the user's move to the board
            user_move_uci = chess.Move.from_uci(user_move)
            if user_move_uci in board.legal_moves:
                board.push(user_move_uci)
                print(f"User move: {user_move}")
            else:
                print("Invalid move! Try again.")
        except ValueError:
            print("Invalid move! Try again.")