# @Author  : lightXu
# @File    : coordinates_correct_pyinstaller.py
# @Time    : 2018/12/10 0010 上午 10:26
import os
import argparse
import cv2
import traceback
import numpy as np
import glob2 as glob
import xml.etree.cElementTree as ET
import requests
import base64
from urllib import parse, request


access_token = '24.214174608e47e6047f31c3fd8c3cedef.2592000.1548390126.282335-14614857'
# access_token = ocr_login()
OCR_BOX_URL = 'https://aip.baidubce.com/rest/2.0/ocr/v1/'
OCR_URL = 'https://aip.baidubce.com/rest/2.0/ocr/v1/'
# OCR_ACCURACY = 'general'
OCR_ACCURACY = 'accurate'
OCR_CLIENT_ID = 'AVH7VGKG8QxoSotp6wG9LyZq'
OCR_CLIENT_SECRET = 'gG7VYvBWLU8Rusnin8cS8Ta4dOckGFl6'
OCR_TOKEN_UPDATE_DATE = 10


def login():
    grant_type = 'client_credentials'
    client_id = OCR_CLIENT_ID
    client_secret = OCR_CLIENT_SECRET

    textmod = {'grant_type': grant_type, 'client_id': client_id, 'client_secret': client_secret}
    textmod = parse.urlencode(textmod)

    # 输出内容:user=admin&password=admin
    header_dict = {'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; Trident/7.0; rv:11.0) like Gecko'}
    url = 'https://aip.baidubce.com/oauth/2.0/token'
    req = request.Request(url='{}{}{}'.format(url, '?', textmod), headers=header_dict)
    res = request.urlopen(req).read()
    token = eval(res.decode(encoding='utf-8'))['access_token']
    return token


def opencv2base64(img):
    image = cv2.imencode('.jpg', img)[1]
    base64_data = str(base64.b64encode(image))[2:-1]
    return base64_data


def get_ocr_raw_result(img, ocr_accuracy=OCR_ACCURACY, language_type='CHN_ENG'):
    textmod = {'access_token': login()}
    textmod = parse.urlencode(textmod)
    url = '{}{}{}{}'.format(OCR_BOX_URL, ocr_accuracy, '?', textmod)
    url_general = '{}{}{}{}'.format(OCR_BOX_URL, 'general', '?', textmod)

    headers = {'Content-Type': 'application/x-www-form-urlencoded'}

    image_type = 'base64'
    group_id = 'group001'
    user_id = 'usr001'

    image = opencv2base64(img)

    data = {
        'image_type': image_type,
        'group_id': group_id,
        'user_id': user_id,
        'image': image,
        'detect_direction': 'true',
        'recognize_granularity': 'small',
        'language_type': language_type,
        # 'vertexes_location': 'true',
        # 'probability': 'true'
    }

    resp = requests.post(url, data=data, headers=headers).json()
    if resp.get('error_msg'):
        if 'internal error' in resp.get('error_msg'):
            resp = requests.post(url_general, data=data, headers=headers).json()
            if resp.get('error_msg'):
                raise Exception("ocr {}!".format(resp.get('error_msg')))
        else:
            raise Exception("ocr {}!".format(resp.get('error_msg')))

    return resp


def get_ocr_text_and_coordinate_in_google_format(img, ocr_accuracy=OCR_ACCURACY, language_type='CHN_ENG'):
    textmod = {'access_token': access_token}
    textmod = parse.urlencode(textmod)
    url = '{}{}{}{}'.format(OCR_BOX_URL, ocr_accuracy, '?', textmod)
    url_general = '{}{}{}{}'.format(OCR_BOX_URL, 'general', '?', textmod)

    headers = {'Content-Type': 'application/x-www-form-urlencoded'}

    image_type = 'base64'
    group_id = 'group001'
    user_id = 'usr001'

    image = opencv2base64(img)

    data = {
        'image_type': image_type,
        'group_id': group_id,
        'user_id': user_id,
        'image': image,
        'detect_direction': 'true',
        'recognize_granularity': 'small',
        'language_type': language_type,
        # 'vertexes_location': 'true',
        # 'probability': 'true'
    }

    resp = requests.post(url, data=data, headers=headers).json()
    if resp.get('error_msg'):
        if 'internal error' in resp.get('error_msg'):
            resp = requests.post(url_general, data=data, headers=headers).json()
            if resp.get('error_msg'):
                raise Exception("ocr {}!".format(resp.get('error_msg')))
        else:
            raise Exception("ocr {}!".format(resp.get('error_msg')))

    words_result = resp.get('words_result')
    dict_list = [item2.get('location') for item in words_result for item2 in item['chars']]
    char_list = [item2.get('char') for item in words_result for item2 in item['chars']]
    matrix = []
    for ele in dict_list:
        xmin = ele['left']
        ymin = ele['top']
        xmax = ele['width'] + ele['left']
        ymax = ele['top'] + ele['height']
        item0 = (xmin, ymin, xmax, ymax)
        matrix.append(item0)

    res_dict = {'chars': char_list, 'coordinates': matrix}
    return res_dict


def read_xml_to_json(xml_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    regions_list = []
    for obj in root.findall('object'):
        class_name = obj.find('name').text
        bbox = obj.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
        bbox_dict = {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}
        region = {'class_name': class_name, 'bounding_box': bbox_dict}
        regions_list.append(region)

    serial = '201812260000001'
    sheet_dict = {'series_number': serial, 'regions': regions_list}
    return sheet_dict


def crop_region(im, bbox):
    xmin = int(bbox['xmin'])
    ymin = int(bbox['ymin'])
    xmax = int(bbox['xmax'])
    ymax = int(bbox['ymax'])

    region = im[ymin:ymax, xmin:xmax]
    return region


def write_single_img(dst, save_path):
    try:
        cv2.imencode('.jpg', dst)[1].tofile(save_path)
    except FileNotFoundError as e:
        raise e


def read_single_img(img_path):
    try:
        im = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
    except FileNotFoundError as e:
        raise e
    return im


def transform(template_img_size, correcting_img, correcting_pts, template_pts):
    pts1 = np.float32(correcting_pts)  # 原始坐标
    pts2 = np.float32(template_pts)  # 目标坐标

    mtx = cv2.getAffineTransform(pts1, pts2)
    dst = cv2.warpAffine(correcting_img, mtx, template_img_size, borderValue=(0, 0, 255))
    return dst


def get_same_str(str1, str2):
    str1_set = set(str1)
    str2_set = set(str2)
    intersection = str1_set & str2_set
    if intersection:
        len1 = len(str1)
        len2 = len(str2)
        if len2 == 1:
            start_index = str1.index(str2)
            return {'ismatch': True, 'coordinates': ((start_index, start_index + 1), (0, 1))}
        else:
            str_set = set(str1 + str2)
            str_set_dict = {}
            for i, ele in enumerate(sorted(list(str_set))):
                str_set_dict[ele] = i + 1

            str1_np = np.asarray([str_set_dict[k] for k in str1])
            str2_np = np.asarray([str_set_dict[k] for k in str2])

            np1 = np.tile(str1_np, (len2, 1))
            np2 = np.tile(str2_np, (1, len1)).reshape(len1, len2).T

            np3 = np1 - np2

            size = np3.shape

            np4 = np3.reshape(-1, 1)

            np4_list = np4.tolist()
            zero_list = list()
            str_index = []
            for i, ele in enumerate(np4):
                if i in zero_list:
                    continue
                else:
                    if ele == [0]:
                        length = 0
                        zero_list.append(i)
                        for interval in range(1, size[1]):
                            next_index = i + size[1] * interval + interval
                            if next_index < len(np4_list):
                                if np4_list[next_index] == [0]:
                                    length += 1
                                    zero_list.append(next_index)  # 跳过的循环
                                if np4_list[next_index] != [0]:
                                    break
                        str_index.append((i, i + size[1] * length + length, length))
                    else:
                        pass

            # print(str_index)
            # print(max_index[0])
            max_index = sorted(str_index, key=lambda k: k[2], reverse=True)
            a = (max_index[0][0] // size[1], max_index[0][0] % size[1])
            b = (max_index[0][1] // size[1], max_index[0][1] % size[1])

            # max_str1 = str1[a[1]:b[1] + 1]
            # max_str2 = str2[a[0]:b[0] + 1]
            # print(max_str1, max_str2)

            return {'ismatch': True, 'coordinates': ((a[1], b[1] + 1), (a[0], b[0] + 1))}

    else:
        return {'ismatch': False, 'coordinates': ()}


def match_string(correcting_ocr, correcting_bias, template_ocr, template_bias):
    correcting_words_list = correcting_ocr['chars']
    template_words_list = template_ocr['chars']

    max_same_str = ''
    index_pair = {}

    longer = ''.join(correcting_words_list)
    shorter = ''.join(template_words_list)

    res = get_same_str(longer, shorter)
    if res['ismatch']:
        c, t = res['coordinates']
        if len(max_same_str) < c[1]-c[0]:
            max_same_str = longer[c[0]:c[1]]
            index_pair['correcting'] = (c[0], c[1])
            index_pair['template'] = (t[0], t[1])

    if len(index_pair) > 0:
        correcting_coordinate = correcting_ocr['coordinates'][index_pair['correcting'][0]]  # xmin, ymin, xmax, ymax
        c_x = correcting_bias[0] + correcting_coordinate[0]
        c_y = correcting_bias[1] + correcting_coordinate[1]
        template_coordinate = template_ocr['coordinates'][index_pair['template'][0]]
        t_x, t_y = template_bias[0] + template_coordinate[0], template_bias[1] + template_coordinate[1]

        print(max_same_str)
        if abs(c_x-t_x) < 50 and abs(c_y-t_y) < 50:
            return {'correcting': (c_x, c_y), 'template': (t_x, t_y)}
        else:
            return {}
    else:
        return {}


def get_template(template_img, ocr_classes_dict):
    template_dict = {}

    for ocr in ocr_classes_dict:
        template = {}
        class_name = ocr['class_name']
        if 'solve' in class_name:
            # ocr['region']['ymax'] = int(0.10 * (ocr['region']['ymax']-ocr['region']['ymin']) +ocr['region']['ymin'])
            ocr['region']['ymax'] = int(250 + ocr['region']['ymin'])
        ocr_box = ocr['region']
        left, top = ocr_box['xmin'], ocr_box['ymin']

        ocr_img = crop_region(template_img, ocr_box)
        # cv2.imshow(class_name, ocr_img)
        # if cv2.waitKey(0) == 27:
        #     cv2.destroyAllWindows()

        ocr_word = get_ocr_text_and_coordinate_in_google_format(ocr_img)  # baidu
        template['words_result'] = ocr_word
        template['coordinate_bias'] = (left, top)

        template_dict[class_name] = template
    return template_dict


def get_correct_points(points_list, size):
    if len(points_list) > size >= 3:
        points_list = sorted(points_list, key=lambda k: k[1])
        i = 1
        choice_index_list = [0]
        for index, ele in enumerate(points_list):
            if abs(ele[1] - points_list[0][1]) > 500:  # 找y轴差值大于500的坐标
                choice_index_list.append(index)
                i = i + 1
                if i == size:
                    break
        return choice_index_list
    elif len(points_list) == 3:
        return [0, 1, 2]
    else:
        raise Exception


def save_transformed_img(template_img_size, images_path_list, ocr_classes_dict,
                         correcting_img_dir_path, corrected_img_save_dir, template_dict):
    for img_path in images_path_list:
        correcting_img = read_single_img(img_path)
        print('***********************************')
        print(img_path)
        correct_coordinates_list = []
        template_coordinates_list = []
        try:
            for ocr in ocr_classes_dict:
                class_name = ocr['class_name']
                if 'solve' in class_name:
                    ocr['region']['ymax'] = ocr['region']['ymax']
                ocr_box = ocr['region']
                left, top = ocr_box['xmin'], ocr_box['ymin']

                correcting_oct_region = crop_region(correcting_img, ocr_box)
                # cv2.imwrite(class_name+'.jpg', correcting_oct_region)
                # cv2.imshow(class_name, correcting_oct_region)
                # if cv2.waitKey(0) == 27:
                #     cv2.destroyAllWindows()
                correcting_word = get_ocr_text_and_coordinate_in_google_format(correcting_oct_region)  # baidu
                template = template_dict[class_name]
                template_word = template['words_result']
                template_bias = template['coordinate_bias']

                coordinates_dict = match_string(correcting_word, (left, top), template_word, template_bias)
                if len(coordinates_dict) > 0:
                    correct_coordinates_list.append(coordinates_dict['correcting'])
                    template_coordinates_list.append(coordinates_dict['template'])
                else:
                    continue

            choice_index_list = get_correct_points(correct_coordinates_list, 3)
            c_coordinates = [correct_coordinates_list[ele] for ele in choice_index_list]
            t_coordinates = [template_coordinates_list[ele] for ele in choice_index_list]
            dst = transform(template_img_size, correcting_img,
                            c_coordinates, t_coordinates)
            print(c_coordinates, t_coordinates)
            save_path = img_path.replace(correcting_img_dir_path, corrected_img_save_dir)
            write_single_img(dst, save_path)
            print(save_path)
        except Exception as e:
            print('image corrected error: {}'.format(e))
            traceback.print_exc()


def correct(template_path, correcting_img_dir_path, corrected_img_save_dir, sheet_dict, sheet_sides='front'):
    find_str = os.path.join(correcting_img_dir_path, '*.jpg')
    correcting_img_path_list = glob.glob(find_str)
    # correcting_img_path_list = [r'C:\Users\Administrator\Desktop\sheet\correct\back_sizes\20180719004308818_0030.jpg']
    template_img = read_single_img(template_path)
    y, x = template_img.shape[0], template_img.shape[1]
    ocr_classes_dict = []
    # ocr_class = ['info_title', 'page']
    ocr_class = {'front': ['info_title', 'page'], 'back': ['solve', 'solve0', 'page']}
    page_index = 1
    for ele in sheet_dict['regions']:
        if ele['class_name'] in ocr_class[sheet_sides]:
            ocr_classes_dict.append({'class_name': '{}_{}'.format(ele['class_name'], str(page_index)),
                                     'region': ele['bounding_box']})
            page_index += 1

    template_info_dict = get_template(template_img, ocr_classes_dict)
    save_transformed_img((x, y), correcting_img_path_list, ocr_classes_dict,
                         correcting_img_dir_path, corrected_img_save_dir, template_info_dict)


def rotate(image, angle, center=None, scale=1.0):
    (h, w) = image.shape[:2]
    if center is None:
        center = (w // 2, h // 2)

    mtx = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, mtx, (w, h))
    return rotated


def image_direction(image_raw, standard_direction):
    # 图片进来取上下两部分，baidu_ocr判断方向
    height, width = image_raw.shape[0], image_raw.shape[1]
    image = ''
    if height > width:
        if standard_direction == 'V':
            image = image_raw
        if standard_direction == 'H':
            image = np.rot90(image_raw)
    if height <= width:
        if standard_direction == 'V':
            image = np.rot90(image_raw)
        if standard_direction == 'H':
            image = image_raw

    crop_ratio = 0.1
    crop_height = int(crop_ratio * height)
    top_part = image[:crop_height, :]
    bottom_part = image[height-crop_height:height, :]
    ocr_used_image = np.vstack([top_part, bottom_part])
    # - -1:未定义，
    # - 0:正向，
    # - 1: 逆时针90度，
    # - 2:逆时针180度，
    # - 3:逆时针270度
    direction = get_ocr_raw_result(ocr_used_image)['direction']
    if direction == 2:
        # flip:  1 水平翻转, 0 垂直翻转, -1 水平垂直翻转
        fliped_image = cv2.flip(image, 1)
    if direction == 0:
        fliped_image = image
    else:
        raise ValueError("direction={} is not supported!".format(direction))

    # 判断是该套试卷中的第几份：识别页码？


def run():
    parser = argparse.ArgumentParser(description="your script description")  # --help
    parser.add_argument('--template', '-t', help='reviewed template file path')
    parser.add_argument('--page', '-p', choices=['front', 'back'], help='front sizes or back sizes')
    parser.add_argument('--raw', '-r', help='raw images folder path')
    parser.add_argument('--save', '-s', help='corrected images save folder path')
    # description参数可以用于插入描述脚本用途的信息，可以为空

    args = parser.parse_args()  # 将变量以标签-值的字典形式存入args字典

    template = args.template
    raw_dir = args.raw
    page = args.page
    save_dir = args.save

    xml_path = template.replace('.jpg', '.xml')
    sheet_dict0 = read_xml_to_json(xml_path)
    correct(template, raw_dir, save_dir, sheet_dict0, sheet_sides=page)


if __name__ == '__main__':
    run()
    # template_path0 = r'C:\Users\Administrator\Desktop\sheet\correct\back_sizes\template\20180719004308818_0020.jpg'
    # img_dir_path = r'C:\Users\Administrator\Desktop\sheet\correct\back_sizes'
    # img_save_dir = r'C:\Users\Administrator\Desktop\sheet\correct\back_sizes\corrected'
    # xml_path = template_path0.replace('.jpg', '.xml')
    # print('hello', xml_path)
    # sheet_dict0 = read_xml_to_json(xml_path)
    # correct(template_path0, img_dir_path, img_save_dir, sheet_dict0, sheet_sides='back')
