import os
import re
from glob import glob

import cv2
import numpy as np
from click_captcha.model import mobile_net, cnn_net, cnn_net_tiny, cnn_net_small
from click_captcha.utils import pil_resize

image_shape = (40, 40, 1)
weights_path = "./models/e01-acc0.83-char.h5"
project_dir = os.path.dirname(os.path.abspath(__file__)) + "/../"


def recognize(image_path):
    model = cnn_net_small(input_shape=image_shape)
    model.load_weights(weights_path)

    paths = glob("../data/test/char_*.jpg")

    X = []
    for image_path in paths:
        print(image_path)
        img = cv2.imread(image_path)
        img = pil_resize(img, image_shape[0], image_shape[1])
        # cv2.imshow("img", img)
        # cv2.waitKey(0)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = np.expand_dims(img, axis=-1)
        img = img / 255.
        # X = np.expand_dims(img, 0)
        X.append(img)

    X = np.array(X)
    print("X.shape", X.shape)
    pred = model.predict(X)
    print("pred.shape", pred.shape)

    with open(project_dir + "data/chinese_6270.txt", 'r') as f:
        char_list = f.readlines()
    char_str = "".join(char_list)
    char_str = re.sub("\n", "", char_str)
    for p in pred:
        index = int(np.argmax(p))
        print(char_str[index], p[index])



    # index_list = []
    # prob_list = []
    # for i in range(5):
    #     index = int(np.argmax(pred))
    #     index_list.append(index)
    #     prob_list.append(np.max(pred))
    #     pred = np.delete(pred, index)
    #
    # index = index_list[0]
    # print("index_list", index_list)
    # print("index", index)
    # with open(project_dir + "data/chinese_6270.txt", 'r') as f:
    #     char_list = f.readlines()
    # char_str = "".join(char_list)
    # char_str = re.sub("\n", "", char_str)
    # char = char_str[index]
    # print("recognize chinese", char, prob_list[0])
    #
    # for i in range(1, len(index_list)):
    #     print("possible chinese", i, char_str[index_list[i]], prob_list[i])
    return


if __name__ == "__main__":
    # _path = "../data/test/char_9.jpg"
    _path = "../data/click/80_70_2.jpg"
    # _path = "../data/click/2019_73_1.jpg"
    recognize(_path)