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, show_cam, calibrate, img_idx = 1, test_img = None, save_img_as = None):
        self.board = chess.Board()
        self.difficulty = difficulty
        self.show_cv = show_cv
        self.show_cam = show_cam
        self.test_img = test_img
        self.calibrate = calibrate
        self.save_img_as = save_img_as
        self.img_idx = int(img_idx)

        self.left_cut = 0
        self.right_cut = 290
        self.top_cut = 0
        self.bottom_cut = 290

        self.picam2 = Picamera2()

    # def calibrate_cam():
    #     pass

    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)

        if self.save_img_as:
            self.calibrate = True

        preview_config = self.picam2.create_preview_configuration(main={"size": (2464, 2464)})
        self.picam2.configure(preview_config)
        if (self.show_cam):
            self.picam2.start_preview(Preview.QTGL)
            # self.picam2.start()
            # input("Press enter to continue.")
        # else:
        #     self.picam2.start()
        self.picam2.start()

        # self.get_img()

        while(1): # essentially do_while(self.calibrate)
            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
            # w_cut = 30
            # h_cut = 100
            cropped_img = orig_img[self.top_cut:h-self.bottom_cut, self.left_cut:w-self.right_cut]
            # cropped_img = orig_img
            img = cv2.resize(cropped_img, (512, 512))
            # img = orig_img

            if (self.show_cv):
                display_img([orig_img, img])
                # display_img([img])

            if (self.calibrate):
                answer = 0
                while(answer != "y" and answer != "n"):
                    answer = input("Done calibrating? (y/n): ")
                    if answer == "y":
                        self.calibrate = False
                    elif answer == "n":
                        self.calibrate = True

            if (not self.calibrate):
                break

        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):
        time.sleep(2)
        if (self.save_img_as):
            img_txt = self.save_img_as + str(self.img_idx) + '.jpg'
        else:
            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)
        self.img_idx += 1
        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("--show_cam", action="store_true", help="Show persistent camera view")
    parser.add_argument("--calibrate", action="store_true", help="Loop before cv")
    parser.add_argument("--img_idx", help="Where to start indexing images (for naming them)")
    parser.add_argument("--test_img", help="If specified, will use said image in test_images folder rather than camera input")
    parser.add_argument("--save_img_as", help="If specified, will save image as given name in game_images")
    args = parser.parse_args()

    game = ChessGame(args.difficulty, args.show_cv, args.show_cam, args.calibrate, args.img_idx, args.test_img, args.save_img_as)
    game.start_game()