game.py

import argparse
import chess
import chess.pgn
from board_detector import find_board
from board_detector import find_pieces
from board_detector import init_show_cv
from board_detector import display_img
import color_analyzer
import move_translator
import cv2
import os

import time
from picamera2 import Picamera2, Preview

class ChessGame:
    def __init__(self, difficulty, show_cv, test_img = None):
        self.board = chess.Board()
        self.difficulty = difficulty
        self.show_cv = show_cv
        self.test_img = test_img
        self.img_idx = 0
        self.picam2 = Picamera2()

    def start_game(self):
        print(f"Starting chess game (difficulty: {self.difficulty})")

        # TODO - call initialize board in board_detector, initialize colors for color analysis,
        # then loop until checkmate. also handle illegal moves (writing to screen if we end up doing that or just LEDs)

        init_show_cv(self.show_cv)

        preview_config = self.picam2.create_preview_configuration(main={"size": (512, 512)})
        self.picam2.configure(preview_config)

        # self.get_img()

        if (self.test_img):
            img_path = os.path.join('test_images', self.test_img)
            orig_img = cv2.imread(img_path)
        else:
            # orig_img = cv2.imread('test_images/board1.jpg') # TODO - CHANGE TO MAKE IT RECEIVE INPUT FROM CAMERA
            orig_img = self.take_pic()
        h,w,c = orig_img.shape
        cropped_img = orig_img[0:h, int(w/2 - h/2 - 60):int(w/2 + h/2 + 100)]
        img = cv2.resize(cropped_img, (512, 512))

        if (self.show_cv):
            # cv2.imshow('Original Image Before Processing', img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()
            display_img([img])

        warped_img, sorted_warped_points = find_board(img)

        if (warped_img is None):
            return
        
        color_grid = find_pieces(warped_img, sorted_warped_points)

        while(1): # game loop
            self.player_turn()
            # handle cheating
            if self.board.is_checkmate():
                break

            quit() # TODO - REMOVE

            self.ai_turn()
            if self.board.is_checkmate():
                break

        # game is over
            
    def player_turn(self):
        # TODO - wait for user button, then check for valid move. loop until a valid move has been made
        pass

    def ai_turn(self):
        # TODO
        # 1. use detect_pieces in color_analyzer
        # 2. update internal representation of board
        # 3. give to python-chess and get best move
        # 4. call move_translator with argument of best move
        pass

    def take_pic(self):
        self.img_idx += 1
        # self.picam2.start_preview(Preview.QTGL)
        self.picam2.start()
        time.sleep(2)
        img_txt = 'board' + str(self.img_idx) + '.jpg'
        print(img_txt)
        img_path = os.path.join('game_images', img_txt)
        metadata = self.picam2.capture_file(img_path)
        return cv2.imread(img_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AI Chess Robot with Computer Vision")
    parser.add_argument("--difficulty", choices=["easy", "medium", "hard"], 
                        default="medium", help="Chess AI difficulty (how far it looks ahead)")
    parser.add_argument("--show_cv", action="store_true", help="Show opencv images as processing occurs during game")
    parser.add_argument("--test_img", help="If specified, will use said image in test_images folder rather than camera input")
    args = parser.parse_args()

    game = ChessGame(args.difficulty, args.show_cv, args.test_img)
    game.start_game()