# @Author  : lightXu
# @File    : ocr_sheet.py
import re
import numpy as np
import xml.etree.cElementTree as ET
from segment.sheet_resolve.tools.utils import create_xml
from segment.sheet_resolve.analysis.sheet.sheet_adjust import adjust_item_edge_by_gray_image


def subfield_answer_sheet(img0, answer_sheet):
    h, w = img0.shape[:2]
    one_part = 0
    line_xmax_1 = 0
    line_xmax_2 = 0
    modules = []
    modules11 = []
    w_int_1 = w
    w_int_2 = round(w / 2)
    w_int_3 = round(w / 3)
    w_int_4 = round(w / 4)
    w_int_8 = round(w / 8)
    if w_int_8 < 50:
        w_int_8 = 50

    key_modules_classes = ['choice', 'cloze', 'solve', 'solve0', 'composition0', 'composition', 'correction',
                           'ban_area', ]
    if h > w:  # 暂定答题卡高大于宽的为单栏
        one_part = 1
    else:
        temp1 = 0
        temp2 = 0
        for ele in answer_sheet:
            if ele["class_name"] in key_modules_classes:
                modules.append(ele)
        modules_xmin = sorted(modules, key=lambda x: (x['bounding_box']['xmin']))
        modules_xmax = sorted(modules, key=lambda x: (x['bounding_box']['xmax']))
        for i in range(len(modules_xmin) - 1):
            if i == 0 and modules_xmin[0]['bounding_box']['xmin'] - 0 > w_int_4:
                temp1 = 1
            else:
                if modules_xmin[i + 1]['bounding_box']['xmin'] - modules_xmin[i]['bounding_box']['xmax'] > w_int_4:
                    if modules11 == []:
                        line_xmax_1 = modules_xmin[i]['bounding_box']['xmax'] + 20
                        line_xmax_2 = modules_xmin[i + 1]['bounding_box']['xmin'] - 20
                    else:
                        modules11.append(modules_xmin[i]['bounding_box']['xmax'])
                        modules11_xmax = sorted(modules11)[-1]
                        line_xmax_1 = modules11_xmax + 20
                        line_xmax_2 = modules_xmin[i + 1]['bounding_box']['xmin'] - 20
                        modules11 = []
                    temp1 = 1
                    temp2 = 1
                    break
                elif modules_xmin[i + 1]['bounding_box']['xmin'] - modules_xmin[i]['bounding_box']['xmax'] > -w_int_8:
                    if temp1 == 0:
                        if modules11 == []:
                            line_xmax_1 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
                                               modules_xmin[i]['bounding_box']['xmax']) / 2)
                        else:
                            modules11.append(modules_xmin[i]['bounding_box']['xmax'])
                            modules11_xmax = sorted(modules11)[-1]
                            line_xmax_1 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
                                               modules11_xmax) / 2)
                            modules11 = []
                        temp1 = 1
                    elif temp1 == 1:
                        if modules11 == []:
                            line_xmax_2 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
                                               modules_xmin[i]['bounding_box']['xmax']) / 2)
                        else:
                            modules11.append(modules_xmin[i]['bounding_box']['xmax'])
                            modules11_xmax = sorted(modules11)[-1]
                            line_xmax_2 = int((modules_xmin[i + 1]['bounding_box']['xmin'] +
                                               modules11_xmax) / 2)
                        temp2 = 1
                else:
                    modules11.append(modules_xmin[i]['bounding_box']['xmax'])

        if temp1 == 0 and temp2 == 0:
            if modules_xmax[-1]['bounding_box']['xmax'] - w < -(2 * w_int_4):
                line_xmax_1 = modules_xmax[-1]['bounding_box']['xmax'] + 20
                line_xmax_2 = 2 * w_int_3
            elif modules_xmax[-1]['bounding_box']['xmax'] - w < -w_int_4:
                line_xmax_1 = modules_xmax[-1]['bounding_box']['xmax'] + 20
        elif temp1 == 1 and temp2 == 0:
            if modules_xmax[-1]['bounding_box']['xmax'] - w < -w_int_4:
                line_xmax_2 = 2 * w_int_3

    return line_xmax_1, line_xmax_2


def tell_columns(image, sheet_dict):
    h, w = image.shape[0], image.shape[1]
    # sheet_dict = adjust_item_edge_by_gray_image(image, sheet_dict)
    x1, x2 = subfield_answer_sheet(image, sheet_dict)

    split_x = [px for px in [x1, x2] if px != 0]

    if not split_x:
        split_x = [w-1]

    return split_x


def ocr2sheet(image, sheet_dict, raw_ocr, xml_path=None):
    image_y, image_x = image.shape[0], image.shape[1]
    col_split_list = tell_columns(image, sheet_dict)
    digital_p = r'\d'
    eng_char_p = '[\u0041-\u005a|\u0061-\u007a]'  # english
    chn_char_p = '[\u4e00-\u9fa5]'  # chinese
    sp_nums_p = '[①②③④⑤⑥⑦⑧⑨⑩]'
    punctuation_p = '[，；：。,;:·√()（）]+'

    pattern_list = [chn_char_p]

    ocr_res_len = len(raw_ocr)
    for i, words_line in enumerate(raw_ocr):
        words = words_line['words']
        words = words.replace(' ', '').upper()  # 去除空格
        loc = words_line['location']
        top = int(loc['top'])
        left = int(loc['left'])
        width = int(loc['width'])
        height = int(loc['height'])
        loc.update({'right': left + width, 'bottom': top + height,
                    'mid_x': left + width // 2, 'mid_y': top + height // 2})

    raw_ocr = sorted(raw_ocr, key=lambda x:x['location']['mid_x'])
    mid_x_list = [ele['location']['mid_x'] for ele in raw_ocr]
    col_list = []
    for split in col_split_list:
        mid_x_list.append(split)
        mid_x_list = sorted(mid_x_list)
        split_index = mid_x_list.index(split)
        col_list.append(raw_ocr[:split_index])
        raw_ocr = raw_ocr[split_index:]
        mid_x_list = mid_x_list[split_index+1:]

    if raw_ocr:
        col_list.append(raw_ocr)

    col_split_list.insert(0, 1)
    col_split_list.append(image_x)

    block_list = []
    for ii, ocr_res in enumerate(col_list):
        ocr_res = sorted(ocr_res, key=lambda x: x['location']['top'])
        raw_chn_index = []
        for i, words_line in enumerate(ocr_res):
            words = words_line['words']
            loc = words_line['location']
            width = int(loc['width'])
            height = int(loc['height'])
            if width >= height:
                match_nums_list = []
                for p in pattern_list:
                    words_m = re.finditer(p, words)
                    match_index_list = [(m.group(), m.span()) for m in words_m if m]
                    match_nums = len(match_index_list) * 2
                    match_nums_list.append(match_nums)

                if sum(match_nums_list) >= 2:
                    raw_chn_index.append(i)

        # print(raw_chn_index)

        left_limit = min([ele['location']['left'] for ele in ocr_res
                          if ele['location']['width'] >= ele['location']['height']]) - 10
        right_limit = max([ele['location']['right'] for ele in ocr_res
                           if ele['location']['width'] >= ele['location']['height']]) + 10

        # 文字识别结果可能连横跨栏，并导致文字分栏错误
        left_limit = max(left_limit, col_split_list[ii])
        right_limit = min(right_limit, col_split_list[ii+1])

        chn_index = raw_chn_index.copy()
        if ocr_res_len - 1 not in raw_chn_index:
            chn_index.append(len(ocr_res) - 1)
        split_index_arr = np.array(chn_index)

        numbers_interval = np.abs(split_index_arr[1:] - split_index_arr[:-1])

        split_index = []
        for i, interval in enumerate(numbers_interval):
            if interval > np.mean(numbers_interval) and interval > 2:
                split_index.append(i)

        split_index = sorted(list(set(split_index)))
        # print('split_index', split_index)

        for i, ele in enumerate(split_index):
            top_limit = raw_chn_index[ele]
            if top_limit == len(ocr_res) - 1:
                break
            else:
                # 下分界行的上一行
                bottom_limit = chn_index[split_index[i] + 1]
                if bottom_limit in raw_chn_index:
                    while int(ocr_res[bottom_limit - 1]["location"]['height']) >= int(
                            ocr_res[bottom_limit - 1]["location"]['width']):
                        bottom_limit = bottom_limit - 1

                    bottom = int(
                        ocr_res[bottom_limit - 1]["location"]["top"] +
                        ocr_res[bottom_limit - 1]["location"]["height"] * 1.2)
                else:
                    bottom_limit = chn_index[-1]
                    bottom = int(
                        ocr_res[bottom_limit]["location"]["top"] +
                        ocr_res[bottom_limit]["location"]["height"] * 1.2)

            # 上分界行的下一行
            top = int(ocr_res[top_limit + 1]["location"]["top"] - 0.2 * ocr_res[top_limit + 1]["location"]["height"])

            left = left_limit
            right = right_limit

            block_list.append({'loc': [left, top, right, bottom]})

    # tree = ET.parse(xml_path)
    # for index, choice_m in enumerate(block_list):
    #     xmin, ymin, xmax, ymax = choice_m["loc"]
    #     tree = create_xml(f'block_{index}', tree, str(xmin), str(ymin), str(xmax), str(ymax))
    #
    # tree.write(xml_path)

    return block_list
