# @Author  : lightXu
# @File    : coordinates_correct.py
# @Time    : 2018/12/10 0010 上午 10:26
import os
import cv2
import traceback
import numpy as np
from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate_in_google_format, tesseract_boxes_by_py
from segment.sheet_resolve.tools.utils import read_single_img, write_single_img, crop_region, read_xml_to_json
import glob2 as glob


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

    M = cv2.getAffineTransform(pts1, pts2)
    dst = cv2.warpAffine(correcting_img, M, 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, c_y = correcting_bias[0] + correcting_coordinate[0], 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, method='google'):
    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()

        if method == 'baidu':
            ocr_word = get_ocr_text_and_coordinate_in_google_format(ocr_img)  # baidu
        else:
            ocr_word = tesseract_boxes_by_py(ocr_img)  # tesseract

        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, method='google'):
    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()
                if method == 'baidu':
                    correcting_word = get_ocr_text_and_coordinate_in_google_format(correcting_oct_region)  # baidu
                else:
                    correcting_word = tesseract_boxes_by_py(correcting_oct_region)  # tesseract

                template = template_dict[class_name]
                template_word = template['words_result']
                template_bias = template['coordinate_bias']

                coordiantes_dict = match_string(correcting_word, (left, top), template_word, template_bias)
                if len(coordiantes_dict) > 0:
                    correct_coordinates_list.append(coordiantes_dict['correcting'])
                    template_coordinates_list.append(coordiantes_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_coordiantes = [template_coordinates_list[ele] for ele in choice_index_list]
            dst = transform(template_img_size, correcting_img,
                            c_coordinates, t_coordiantes)
            print(c_coordinates, t_coordiantes)
            save_path = img_path.replace(correcting_img_dir_path, corrected_img_save_dir)
            write_single_img(dst, save_path)
            print(save_path)
        except Exception:
            print('image corrected error')
            traceback.print_exc()


def correct(template_path, correcting_img_dir_path, corrected_img_save_dir, sheet_dict, method, 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, method)
    save_transformed_img((x, y), correcting_img_path_list, ocr_classes_dict, correcting_img_dir_path, corrected_img_save_dir, template_info_dict, method)


if __name__ == '__main__':
    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')
    sheet_dict0 = read_xml_to_json(xml_path)
    correct(template_path0, img_dir_path, img_save_dir, sheet_dict0, sheet_sides='back', method='baidu')
