# @Author  : lightXu
# @File    : mark_box.py
# @Time    : 2018/11/21 0021 下午 16:16
import time
import re
import cv2
import xml.etree.cElementTree as ET

from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate
from segment.sheet_resolve.tools import utils


def solve_mark(left, top, solve_img, xml_path):
    shape = solve_img.shape
    y, x = shape[0], shape[1]

    # ocr_region = solve_img[0:int(0.15 * y), :]
    ocr_region = solve_img[0:250, :]

    # cv2.imshow('ocr_region', ocr_region)
    # if cv2.waitKey(0) == 27:
    #     cv2.destroyAllWindows()

    t11 = time.time()
    word_result_list = get_ocr_text_and_coordinate(ocr_region)
    t22 = time.time()
    print('mark ocr time cost: ', t22-t11)

    if len(word_result_list) < 1:
        return {}
    else:
        words_str = str([ele['words'] for ele in word_result_list])
        number = 999
        number_model = re.compile("\d+[.、:：]\D")
        number_list = number_model.findall(words_str)
        if len(number_list) > 0:
            number = int(re.sub('[\D]', '', number_list[0]))

        all_chars_list = []
        zhmodel = re.compile(u'[\u4e00-\u9fa5]')

        for i, chars_dict in enumerate(word_result_list):
            words = re.sub('[iIl|点]', '1', chars_dict['words'])
            match = zhmodel.search(words)  # 是否有中文
            if not match:
                chars_list = chars_dict['chars']
                all_chars_list = all_chars_list + chars_list

        new_all_chars_list = []
        i = 1

        while i <= len(all_chars_list):
            pre_one = all_chars_list[i - 1]
            if i == len(all_chars_list):
                new_all_chars_list.append(pre_one)
                break
            rear_one = all_chars_list[i]
            condition1 = abs(pre_one['location']['top'] - rear_one['location']['top']) < pre_one['location'][
                'height']  # 两字高度差小于一字高度
            condition2 = pre_one['location']['left'] + 2 * pre_one['location']['width'] > rear_one['location'][
                'left']  # 两字长度大于两字间间隔
            if condition1:
                if condition2:
                    new_char = pre_one['char'] + rear_one['char']
                    new_location = {'left': pre_one['location']['left'],
                                    'top': min(pre_one['location']['top'], rear_one['location']['top']),
                                    'width': rear_one['location']['left'] + rear_one['location']['width'] -
                                    pre_one['location']['left'],
                                    'height': max(pre_one['location']['height'], rear_one['location']['height'])}
                    new_all_chars_list.append({'char': new_char, 'location': new_location})
                    i = i + 1 + 1
                else:
                    new_all_chars_list.append(pre_one)
                    i = i + 1
            else:
                new_all_chars_list.append(pre_one)  # 遇到字符y轴相差过大就结束
                break  # break 直接跳行

        tree = ET.parse(xml_path)

        xml_list = []
        for i, ele in enumerate(new_all_chars_list[1:]):  # 从第二位开始索引
            pre_one = new_all_chars_list[i]
            rear_one = ele

            intervel = (rear_one['location']['left'] -
                        pre_one['location']['left'] -
                        pre_one['location']['width']) // 2

            xmin = ele['location']['left'] - intervel + left
            xmax = xmin + ele['location']['width'] + 2 * intervel
            ymin = ele['location']['top'] - ele['location']['height'] // 2 + top
            ymax = ymin + 2 * ele['location']['height']
            tree = utils.create_xml('{}_solve{}'.format(number, ele['char']), tree, xmin, ymin, xmax, ymax)
            xml_list.append({'char': ele['char'], 'location': [xmin, ymin, xmax, ymax]})

            if i == 0:
                intervel = (rear_one['location']['left'] -
                            pre_one['location']['left'] -
                            pre_one['location']['width']) // 2

                xmin = pre_one['location']['left'] - intervel + left
                xmax = xmin + pre_one['location']['width'] + 2 * intervel
                ymin = pre_one['location']['top'] - pre_one['location']['height'] // 2 + top
                ymax = ymin + 2 * pre_one['location']['height']
                tree = utils.create_xml('{}_solve{}'.format(number, pre_one['char']), tree, xmin, ymin, xmax, ymax)
                xml_list.insert(0, {'char': pre_one['char'], 'location': [xmin, ymin, xmax, ymax]})

        tree.write(xml_path)
        return {'number': number, 'mark': xml_list}


if __name__ == '__main__':
    path = r'C:\Users\Administrator\Desktop\sheet\cloze01.jpg'
    image = cv2.imread(path)
    t1 = time.time()
    # solve_mark(image, path)
    t2 = time.time()
    print(t2 - t1)
