import cv2
import itertools, os, time
import numpy as np
from Model import get_Model
from parameter import letters
import argparse
from keras import backend as K
K.set_learning_phase(0)

Region = {"A": "서울 ", "B": "경기 ", "C": "인천 ", "D": "강원 ", "E": "충남 ", "F": "대전 ",
          "G": "충북 ", "H": "부산 ", "I": "울산 ", "J": "대구 ", "K": "경북 ", "L": "경남 ",
          "M": "전남 ", "N": "광주 ", "O": "전북 ", "P": "제주 "}
Hangul = {"dk": "아", "dj": "어", "dh": "오", "dn": "우", "qk": "바", "qj": "버", "qh": "보", "qn": "부",
          "ek": "다", "ej": "더", "eh": "도", "en": "두", "rk": "가", "rj": "거", "rh": "고", "rn": "구",
          "wk": "자", "wj": "저", "wh": "조", "wn": "주", "ak": "마", "aj": "머", "ah": "모", "an": "무",
          "sk": "나", "sj": "너", "sh": "노", "sn": "누", "fk": "라", "fj": "러", "fh": "로", "fn": "루",
          "tk": "사", "tj": "서", "th": "소", "tn": "수", "gj": "허"}

def decode_label(out):
    # out : (1, 32, 42)
    out_best = list(np.argmax(out[0, 2:], axis=1))  # get max index -> len = 32
    out_best = [k for k, g in itertools.groupby(out_best)]  # remove overlap value
    outstr = ''
    for i in out_best:
        if i < len(letters):
            outstr += letters[i]
    return outstr


def label_to_hangul(label):  # eng -> hangul
    region = label[0]
    two_num = label[1:3]
    hangul = label[3:5]
    four_num = label[5:]

    try:
        region = Region[region] if region != 'Z' else ''
    except:
        pass
    try:
        hangul = Hangul[hangul]
    except:
        pass
    return region + two_num + hangul + four_num

parser = argparse.ArgumentParser()
parser.add_argument("-w", "--weight", help="weight file directory",
                    type=str, default="models/weights.best.hdf5")
parser.add_argument("-t", "--test_img", help="Test image directory",
                    type=str, default="./DB/test/")
args = parser.parse_args()

# Get CRNN model
model = get_Model(training=False)

try:
    model.load_weights(args.weight)
    print("...Previous weight data...")
except:
    raise Exception("No weight file!")


test_dir =args.test_img
test_imgs = os.listdir(args.test_img)
total = 0
acc = 0
letter_total = 0
letter_acc = 0
start = time.time()
for test_img in test_imgs:
    img = cv2.imread(test_dir + test_img, cv2.IMREAD_GRAYSCALE)

    img_pred = img.astype(np.float32)
    img_pred = cv2.resize(img_pred, (128, 64))
    img_pred = (img_pred / 255.0) * 2.0 - 1.0
    img_pred = img_pred.T
    img_pred = np.expand_dims(img_pred, axis=-1)
    img_pred = np.expand_dims(img_pred, axis=0)

    net_out_value = model.predict(img_pred)

    pred_texts = decode_label(net_out_value)

    for i in range(min(len(pred_texts), len(test_img[0:-4]))):
        if pred_texts[i] == test_img[i]:
            letter_acc += 1
    letter_total += max(len(pred_texts), len(test_img[0:-4]))
    predOk = "True"
    if pred_texts == test_img[0:-4]:
        acc += 1
    else:
        predOk = "False"
    total += 1
    # print('Predicted: %s  /  True: %s / net_out_value: %s / ' % (label_to_hangul(pred_texts), label_to_hangul(test_img[0:-4])))
    print('Predicted: %s  /  True: %s / predOk: %s ' % (pred_texts, test_img[0:-4], predOk ))
    
    # cv2.rectangle(img, (0,0), (150, 30), (0,0,0), -1)
    # cv2.putText(img, pred_texts, (5, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255),2)

    #cv2.imshow("q", img)
    #if cv2.waitKey(0) == 27:
    #   break
    #cv2.destroyAllWindows()

end = time.time()
total_time = (end - start)
print("Time : ",total_time / total)
print("ACC : ", acc / total)
print("letter ACC : ", letter_acc / letter_total)

[((8, 0, 8, 23), 'keras.backend.set_learning_phase', 'K.set_learning_phase', ({(8, 21, 8, 22): '(0)'}, {}), '(0)', True, 'from keras import backend as K\n'), ((46, 9, 46, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((54, 8, 54, 33), 'Model.get_Model', 'get_Model', (), '', False, 'from Model import get_Model\n'), ((64, 12, 64, 37), 'os.listdir', 'os.listdir', ({(64, 23, 64, 36): 'args.test_img'}, {}), '(args.test_img)', False, 'import itertools, os, time\n'), ((69, 8, 69, 19), 'time.time', 'time.time', ({}, {}), '()', False, 'import itertools, os, time\n'), ((105, 6, 105, 17), 'time.time', 'time.time', ({}, {}), '()', False, 'import itertools, os, time\n'), ((71, 10, 71, 63), 'cv2.imread', 'cv2.imread', ({(71, 21, 71, 40): 'test_dir + test_img', (71, 42, 71, 62): 'cv2.IMREAD_GRAYSCALE'}, {}), '(test_dir + test_img, cv2.IMREAD_GRAYSCALE)', False, 'import cv2\n'), ((74, 15, 74, 46), 'cv2.resize', 'cv2.resize', ({(74, 26, 74, 34): 'img_pred', (74, 36, 74, 45): '(128, 64)'}, {}), '(img_pred, (128, 64))', False, 'import cv2\n'), ((77, 15, 77, 48), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((78, 15, 78, 47), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((21, 20, 21, 49), 'numpy.argmax', 'np.argmax', (), '', True, 'import numpy as np\n'), ((22, 30, 22, 57), 'itertools.groupby', 'itertools.groupby', ({(22, 48, 22, 56): 'out_best'}, {}), '(out_best)', False, 'import itertools, os, time\n')]