import numpy as np
from keras.models import Model, load_model
from keras.utils import to_categorical
import cv2
from skimage import color, exposure
from _cutter import image_cut

def man_result_check():
    pred_y = np.load("predicted_results.npy")
    test_y = np.load("Waldo_test_lbl.npy")

    test_y = to_categorical(test_y)

    f = open("test_output.txt", 'w')
    z = 0
    for i in range(0, len(test_y)):
        print(pred_y[i], test_y[i], file=f)
        # Calculates correct predictions
        if pred_y[i][0] == test_y[i][0]:
            z+=1

    print("Accuracy: {}%".format(z/len(test_y)*100))
    f.close()

'''
Purpose:Loads a trained neural network model (using Keras) to classify an image
Input:  path/to/trained_model
        image [or] path/to/image [if from_file=True]
Returns:Boolean variable
'''
def is_Wally(trained_model_path, image, from_file=False):
    if from_file:
        img = cv2.imread(image)         # Opens the image (in BGR format)

        # Histogram normalization in v channel
        hsv = color.rgb2hsv(img)
        hsv[:, :, 2] = exposure.equalize_hist(hsv[:, :, 2])
        img = color.hsv2rgb(hsv)

        image = np.rollaxis(img, -1)      # Rolls the colour axis to the front

    trained_model = load_model(trained_model_path)
    if trained_model.predict(image, verbose=1, batch_size=1)[0] == 1:
        return 0
    else:
        return 1

# Load fully puzzle image
# Split image into array of images
# use is_Wally(img) to classify image
# Mark Wally image somehow (colour the border)
# Stitch original image back together

if __name__ == '__main__':
    # Read image
    image = cv2.imread("10.jpg")
    # Split image
    cuts = image_cut(image, 64, 64)
    for i in range(len(cuts)):
        # Transform block
        hsv = color.rgb2hsv(cuts[i])
        hsv[:, :, 2] = exposure.equalize_hist(hsv[:, :, 2])
        block = color.hsv2rgb(hsv)
        block = np.rollaxis(block, -1)
        if is_Wally("Waldo.h5", block):
        # if True:
            # Border block
            GREEN = [0, 255, 0]
            cuts[i] = cv2.copyMakeBorder(cuts[i][1:61,1:61],2,2,2,2,cv2.BORDER_CONSTANT,value=GREEN)

    # Stitch image TODO!
    # Show image
    cv2.imwrite('output.png',image)