import cv2
import numpy as np
from tensorflow.python.keras import models

from util import PIECE, Squares, Board

new_model = models.load_model('chess_model_3_pieces.h5')
#new_model.summary()

#board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
#board = cv2.imread("whole_boards/boards_for_empty/board_1554285167.655788_rank_5.png")
board = cv2.imread("whole_boards/boards_for_empty/board_1554288891.129901_rank_8.png")


def predict_board(occupied_squares: Squares) -> Board:
    board = Board()
    for pos, square in occupied_squares.items():
        square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
        width, height = square.shape
        square = square / 255.0
        test = new_model.predict(np.array(square).reshape((-1, width, height, 1)))

        #cv2.putText(square, f"{pos} {PIECE(int(np.argmax(test)))}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255,) * 3, thickness=3)
        #cv2.imwrite("lel", square)
        board[pos] = PIECE(int(np.argmax(test)))

    return board