import cv2
import sys
from collections import defaultdict
from datetime import datetime

import numpy as np

import runner
from util import load_classifier, PIECE, COLOR, POSITION, Board, Squares, PieceAndColor

np.set_printoptions(threshold=sys.maxsize)


def identify_piece(image: np.ndarray, sift : cv2.xfeatures2d_SIFT, empty_bias=False) -> PieceAndColor:
    centers = np.load("training_data/centers.npy")
    probs = defaultdict(lambda: defaultdict(float))
    best = 0
    best_piece = best_color = None
    for piece in PIECE:
        for color in COLOR:
            #color = runner.compute_color(file, rank)
            classifier = load_classifier(f"classifiers/classifier_{piece}/{color}.pkl")
            features = runner.generate_bag_of_words(image, centers, sift)
            prob = classifier.predict_proba(features)
            probs[piece][color] = prob[0, 1]
            if prob[0, 1] > best:
                best_piece, best_color = piece, color
    print(probs)

    if empty_bias:
        probs[PIECE.EMPTY] *= 1.2

    return best_piece, best_color


def pred_test(position: POSITION, mystery_image=None, empty_bias=False):
    sift = cv2.xfeatures2d.SIFT_create()
    if mystery_image is None:
        mystery_image = cv2.imread("training_images/rook/white/rook_training_D4_2.png")
    probs = classify(mystery_image, sift, empty_bias=empty_bias)
    return probs


def pre_process_and_train() -> None:
    runner.do_pre_processing()
    runner.train_pieces_svm()


def build_board_from_squares(squares: Squares) -> Board:
    sift = cv2.xfeatures2d.SIFT_create()
    board = Board()
    counter = 0
    for position, square in squares.values():
        likely_piece = identify_piece(square, sift)
        board[position] = likely_piece
        if likely_piece != PIECE.EMPTY:
            counter += 1

    print(counter)
    print(64/(counter-1))
    return board


def test_entire_board() -> None:
    board_img = cv2.imread("homo_pls_fuck.jpg")
    warped = runner.warp_board(board_img)

    squares = runner.get_squares(warped)
    board = build_board_from_squares(squares)
    print(board)


def predict(square: np.ndarray, position: POSITION) -> PIECE:
    y, x = np.histogram(square.ravel(), bins=256, range=[0, 256])

    for color in COLOR:
        empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_{color}_square.pkl")
        prob = empty_classifier.predict_proba(np.array(y).reshape(1, -1))
        print(f"{file}{rank}, {color}: {prob[0, 1]}")
        if prob[0, 1] > 0.5:
            return PIECE.EMPTY

    return None


if __name__ == '__main__':
    board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
    warped = runner.warp_board(board)

    empty = 0

    files = "ABCDEFGH"
    ranks = [1, 2, 3, 4, 5, 6, 7, 8]


    non_empties = []

    for position in POSITION:
        counter = 0

        src = runner.get_square(warped, position)
        width, height, _ = src.shape
        src = src[width//25:, height//25:]
        # src = src[:-width//200, :-height//200]
        segmentator = cv2.ximgproc.segmentation.createGraphSegmentation(sigma=0.8, k=150, min_size=700)
        segment = segmentator.processImage(src)

        mask = segment.reshape(list(segment.shape) + [1]).repeat(3, axis=2)
        masked = np.ma.masked_array(src, fill_value=0)
        pls = []
        for i in range(np.max(segment)):
            masked.mask = mask != i

            y, x = np.where(segment == i)

            top, bottom, left, right = min(y), max(y), min(x), max(x)

            dst = masked.filled()[top: bottom + 1, left: right + 1]
            lel = (bottom - top) * (right - left)
            #print(f"this is lel: {lel} ")
            #print(f"this is meh: {np.sum(mask[:,:,0])} ")
            if position == POSITION.H7:
                print("--"*20)
                print("H7")
                print(lel)
                print(len(y))
                print(np.max(segment))
#                    print(lel)
#                    print(np.sum(mask[:, :, 0]))
                print("--"*20)


            pls.append(len(y))
            if len(y) < (164**2)*0.65:
                counter += 1

            cv2.imwrite(f"segment_test/segment_{datetime.utcnow().timestamp()}_{position}.png", dst)

        if np.max(segment) > 0 and not np.all([x < (164**2)*0.2 for x in pls]) and (np.max(segment) >= 3 or np.all([x < (164**2)*0.942 for x in pls])):
            print(f"{position} is nonempty")
            non_empties.append([f"{position}", src])
            print(counter)
            print(np.max(segment))
            empty += 1
    print("++"*20)
    print(counter)
    print(64-empty)

    for non_empty in non_empties:
        cv2.imshow(non_empty[0], non_empty[1])
    cv2.waitKey(0)
    exit()

    #empty_classifier = load_classifier(f"classifiers/classifier_empty/white_piece_on_white_square.pkl")
    #print(empty_classifier.predict_proba(np.array([0]*16).reshape(1, -1))[0, 1])

    #exit()


    board = cv2.imread("whole_boards/board_102_1554110461.608167_.png")
    warped = runner.warp_board(board)

    files = "ABCDEFGH"
    ranks = [1, 2, 3, 4, 5, 6, 7, 8]

    counter = 0

    for file in files:
        for rank in ranks:
            square_img = runner.get_square(warped, file, rank)
            if predict(square_img, file, rank) == 'empty':
                counter += 1

    print(counter)
    exit()


    square_img = runner.get_square(warped, "D", 2)

    gray_square_img = cv2.cvtColor(square_img, cv2.COLOR_BGR2GRAY)
    print(cv2.meanStdDev(gray_square_img)[1])
    print(cv2.meanStdDev(square_img)[1])
    cv2.imshow("square", square_img)
    cv2.waitKey(0)




    print(pred_test("C", 2, square_img))

    sift: cv2.xfeatures2d_SIFT = cv2.xfeatures2d.SIFT_create()
    gray = cv2.cvtColor(square_img, cv2.COLOR_BGR2GRAY)

    kp, desc = sift.detectAndCompute(gray, None)

    cv2.drawKeypoints(square_img, kp, square_img)

    cv2.imshow("kp", square_img)
    cv2.waitKey(0)