# @Author  : lightXu
# @File    : cloze_box.py
# @Time    : 2018/11/23 0023 上午 10:47
import re
import xml.etree.cElementTree as ET
import time
from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate
from segment.sheet_resolve.tools import tf_settings, utils


def get_cloze_box_coordinate(solve_img):
    ocr_region = solve_img
    t11 = time.time()
    word_result_list = get_ocr_text_and_coordinate(ocr_region)
    t22 = time.time()
    print('choice ocr time cost: ', t22-t11)
    return word_result_list


def get_each_coordinate(solve_img):

    word_result_list = get_cloze_box_coordinate(solve_img)
    all_char_list = []
    digital_model = re.compile(r'\d')
    for i, chars_dict in enumerate(word_result_list):
        chars_list = chars_dict['chars']
        for ele in chars_list:
            if digital_model.search(ele['char']):
                all_char_list.append(ele)

    return all_char_list


def decide_coordinate_contains(solve_img, xml_box):
    ocr_box = get_each_coordinate(solve_img)

    ocr_all_char_list = []
    i = 1
    while i <= len(ocr_box):
        pre_one = ocr_box[i - 1]
        if i == len(ocr_box):
            ocr_all_char_list.append(pre_one)
            break
        rear_one = ocr_box[i]
        # 两字高度差小于一字高度
        condition1 = abs(pre_one['location']['top'] - rear_one['location']['top']) < pre_one['location']['height']
        # 两字长度大于两字间间隔
        condition2 = abs(pre_one['location']['left'] + pre_one['location']['width']
                         - rear_one['location']['left']) < pre_one['location']['width']
        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'])}
                ocr_all_char_list.append({'char': new_char, 'location': new_location})
                i = i + 1 + 1
            else:
                ocr_all_char_list.append(pre_one)
                i = i + 1
        else:
            ocr_all_char_list.append(pre_one)  # 遇到字符y轴相差过大就结束
            i = i + 1

    content_list = []
    for index, xml_b in enumerate(xml_box):  # faster-rcnn 的框
        bbox_right = []
        xml_b = xml_b['bounding_box']
        xmin2 = xml_b['xmin']
        ymin2 = xml_b['ymin']
        xmax2 = xml_b['xmax']
        ymax2 = xml_b['ymax']
        mid_x2 = int(xmin2 + (xmax2 - xmin2) // 2)
        box_coordiante = (xmin2, ymin2, xmax2, ymax2)
        choice_number = {'number': 999, 'location': box_coordiante}
        content_list.insert(index, choice_number)

        fixed_height = 60  # 高度固定
        for ocr_b in ocr_all_char_list:  # ocr识别的框
            xmin1 = ocr_b['location']['left']
            ymin1 = ocr_b['location']['top']
            xmax1 = xmin1 + ocr_b['location']['width']
            ymax1 = ymin1 + ocr_b['location']['height']
            mid_x = int(xmin1 + (xmax1 - xmin1) // 2)
            mid_y = int(ymin1 + (ymax1 - ymin1) // 2)

            if xmin2 <= mid_x <= xmax2 and ymin2 <= mid_y <= ymax2 and mid_x < mid_x2:  # 包含且在左侧
                content_list[index]['number'] = ocr_b['char']

            if xmin2 <= mid_x <= xmax2 and ymin2 <= mid_y <= ymax2 and mid_x >= mid_x2:  # 包含且在右侧
                ocr_b['location']['left'] = int(xmin1 - 1 * ocr_b['location']['width'])  # 打分框的边框
                ocr_b['location']['width'] = 3 * ocr_b['location']['width']
                ocr_b['location']['top'] = int(ymin1 - 0.5 * ocr_b['location']['height'])

                xmin = ocr_b['location']['left']
                xmax = xmin + ocr_b['location']['width']
                ymin = ocr_b['location']['top']
                ymax = ymin + fixed_height
                bbox_right.append({'points': ocr_b['char'], 'location': [xmin, ymin, xmax, ymax]})

        ymin = min([ele['location'][1] for ele in bbox_right])
        ymax = max([ele['location'][3] for ele in bbox_right])
        for ele in bbox_right:
            ele['location'][1] = ymin
            ele['location'][3] = ymax

        content_list[index]['right'] = bbox_right

    return content_list


def cloze(left, top, image, choice_bbox_list, xml_path):

    content = decide_coordinate_contains(image, choice_bbox_list)

    tree = ET.parse(xml_path)  # xml tree
    cloze_list = []
    for ele in content:
        number = '{:02d}_cloze'.format(int(ele['number']))
        cloze_xmin = ele['location'][0] + left
        cloze_ymin = ele['location'][1] + top
        cloze_xmax = ele['location'][2] + left
        cloze_ymax = ele['location'][3] + top

        tree = utils.create_xml(number, tree, cloze_xmin, cloze_ymin, cloze_xmax, cloze_ymax)
        region = [cloze_xmin, cloze_ymin, cloze_xmax, cloze_ymax]

        points_list = []
        if len(ele['right']) > 0:  # 存在打分框
            for right_ele in ele['right']:
                points = right_ele['points']
                xmin = right_ele['location'][0] + left
                ymin = right_ele['location'][1] + top
                xmax = right_ele['location'][2] + left
                ymax = right_ele['location'][3] + top
                point_number = '{}_{}'.format(number, points)

                tree = utils.create_xml(point_number, tree, xmin, ymin, xmax, ymax)
                region = [xmin, ymin, xmax, ymax]
                points_list.append({'points': point_number, 'region': region})

        cloze_list.append({'number': number, 'region': region, 'points': points_list})
    tree.write(xml_path)
    return cloze_list
