import time

from models.detection.detector import *
from models.detection.board_manager import *
from models.detection.pieces_manager import *

def print_board_console(board_array):
    files = "abcdefgh"
    print("  " + " ".join(files))
    print(" +-----------------+")

    for r in range(8):
        row_str = str(8 - r) + "|"
        for f in range(8):
            piece = board_array[r][f]
            if piece is None:
                row_str += ". "
            else:
                if piece[0] == "w":
                    piece = piece.upper()
                row_str += piece[2] + " "
        row_str += "|"
        print(row_str)
    print(" +-----------------+")

if __name__ == "__main__":

    print("Initializing models...")
    edges_detector = Detector("../assets/models/edges.pt")
    pieces_detector = Detector("../assets/models/unified-nano-refined.pt")

    pieces_manager = PiecesManager()
    board_manager = BoardManager()

    print("Initializing camera...")
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    cap.set(cv2.CAP_PROP_FPS, 30)

    print("Initialized")
    if not cap.isOpened():
        print("Error: Could not open camera")
        exit()

    cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)

    while True:

        ret, frame = cap.read()
        if not ret:
            print("Error: Failed to grab frame")
            break

        edges_pred = edges_detector.make_prediction(frame)
        pieces_pred = pieces_detector.make_prediction(frame)

        edges_annotated_frame = edges_pred[0].plot()
        pieces_annotated_frame = pieces_pred[0].plot(img=edges_annotated_frame)

        cv2.imshow("Predictions", pieces_annotated_frame)
        cv2.resizeWindow("Predictions", 640, 640)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        remap_width = 800
        remap_height = 800

        board_manager = BoardManager()
        result = board_manager.process_frame(edges_pred[0], frame, (remap_width, remap_height))
        if result is None:
            continue

        corners = result[0]
        matrix = result[1]

        detections = pieces_manager.extract_pieces(pieces_pred)

        board = pieces_manager.pieces_to_board(detections, matrix, (remap_width, remap_height))

        fen = pieces_manager.board_to_fen(board)
        print("FEN:", fen)

        print_board_console(board)

    cap.release()
    cv2.destroyAllWindows()