# @Author  : lightXu
# @File    : sheet_points.py
# @Time    : 2019/9/12 0012 下午 14:39
import re
import numpy as np
from segment.sheet_resolve.tools import utils
from segment.sheet_resolve.tools.brain_api import get_ocr_text_and_coordinate
OCR_ACCURACY = 'accurate'


def split_col(box_list, width):
    if len(box_list) == 1:
        return [box_list]
    else:
        box_list = sorted(box_list, key=lambda k: k.get('bounding_box')['xmin'])
        box_list_class_name = [[ele_box['bounding_box']['xmin'], ele_box['bounding_box']['ymin'],
                                ele_box['bounding_box']['xmax'], ele_box['bounding_box']['ymax']] for ele_box in
                               box_list]
        box_list_raw = sorted(box_list_class_name, key=lambda k: k[0])
        box_list_raw_arr = np.array(box_list_raw)

        pre = box_list_raw_arr[1:, 0]
        rear = box_list_raw_arr[:-1, 0]
        y_diff = pre - rear
        index_list = [index for index, ele in enumerate(y_diff) if abs(ele) > width // 10]

        res_list = []
        if index_list == []:
            res_list = [sorted(box_list, key=lambda k: k.get('bounding_box')['ymin'])]
            return res_list
        else:
            split_x_index = [ele + 1 for ele in index_list]
            split_x_index.insert(0, 0)
            split_x_index.insert(-1, len(box_list_raw))
            split_x_index = sorted(list(set(split_x_index)))
            for i, split in enumerate(split_x_index[1:]):
                one_col = box_list[split_x_index[i]:split_x_index[i + 1]]
                one_col = sorted(one_col, key=lambda k: k.get('bounding_box')['ymin'])
                res_list.append(one_col)
        return res_list


def check_classes(sheet_dict_of_front):
    sheet_dict_of_front_temp = sheet_dict_of_front.copy()
    index_list = []
    for index, ele in enumerate(sheet_dict_of_front_temp):
        for index0, one_paper in enumerate(ele):
            if one_paper['title_number'] == -1 and one_paper['title_with_value'] != -1:
                index_list.append(index0)
    # print(index_list)
    if index_list == []:
        return sheet_dict_of_front
    else:
        if 0 in index_list and len(index_list) > 2:
            index_list = index_list[1:]
        elif 0 in index_list and len(index_list) == 1:
            return sheet_dict_of_front
        else:
            index_list = index_list
        for index_a, ele_a in enumerate(sheet_dict_of_front):
            for index1, ele1 in enumerate(index_list):
                former_class = ele_a[ele1 - 1]
                if former_class != []:
                    lack_title_number_dict = ele_a[ele1]
                    lack_title_number_dict['title_number'] = former_class['title_number']
                    lack_title_number_dict.update({'title_number': former_class['title_number']})
        return sheet_dict_of_front


def change_box(cloze_s_res, cloze_s_region):
    words_result = []
    for ele in cloze_s_res:
        location = ele['location']
        xmin = location['left']
        ymin = location['top']
        xmax = location['left'] + location['width']
        ymax = location['top'] + location['height']
        bbox0 = utils.get_img_region_box01([xmin, ymin, xmax, ymax], cloze_s_region)
        location.update({'left': bbox0[0], 'top': bbox0[1], 'width': bbox0[2] - bbox0[0], 'height': bbox0[3] - bbox0[1]})
        words_result.append(ele)
    return words_result


def get_total_title_quantity_and_value(box_with_content):
    list_of_all = []

    for index, ele in enumerate(box_with_content):
        total_score = []
        title_number = []

        title_two_number = []
        title_two_value = []
        title_two_number_value_list = []

        title_three_number = []
        title_three_value = []

        above_content = ele['above_content']
        for words_index, words_str in enumerate(above_content):
            words = words_str['words']

            pattern1 = re.compile(
                '^\d+[,、.]?[\u4e00-\u9fa5]?[(（]?\d+分+[)）]?[(（]?\d+[)）]?|^\d+[,、.]?[\u4e00-\u9fa5]?[(（]?\d+分+[)）]?')
            result1 = re.findall(pattern1, words)

            pattern11 = re.compile('^\d+[,、.]?[\u4e00-\u9fa5]?[(（]?\d+[)）]?')
            result11 = re.findall(pattern11, words)

            pattern2 = re.compile('[(（]?\d+[)）]?')
            result2 = re.findall(pattern2, words)

            pattern3 = re.compile('[(（]?\d?分+[)）]?')
            result3 = re.findall(pattern3, words)

            pattern4 = re.compile(r'①|②|③|④|⑤|⑥|⑦|⑧|⑨|⑩')
            result4 = re.findall(pattern4, words)

            pattern5 = re.compile('[\u4e00-\u9fa5]')
            result5 = re.findall(pattern5, words)

            if result1 and title_number == []:
                title_number_and_value = ''.join(result1)
                digital_number = re.findall('\d+', title_number_and_value)
                if len(digital_number) == 1:
                    title_number.append(int(digital_number[0]))
                    continue
                elif len(digital_number) == 2:
                    title_number.append(int(digital_number[0]))
                    total_score.append(int(digital_number[1]))
                    continue
                elif len(digital_number) == 3:
                    title_number.append(int(digital_number[0]))
                    total_score.append(int(digital_number[1]))
                    title_two_value.append(int(digital_number[2]))
                    continue

            if result11 and result2 and not result3:
                digital_number = re.findall('\d+', words)
                if len(digital_number) == 1:
                    title_number.append(int(digital_number[0]))
                else:
                    title_number.append(int(digital_number[0]))
                    title_two_number.append(int(digital_number[1]))

            if result2 and not result1 and not result3 and not result11:
                title_two_number_group = result2[0]
                digital_number2 = re.search('\d+', title_two_number_group)
                digital_number22 = digital_number2.group()
                title_two_number.append(int(digital_number22))

            if result3 and not result1 and not result2 and not result4 and not result5:
                title_two_value_group = result3[0]
                digital_value = re.search('\d+', title_two_value_group)
                if digital_value == None:
                    value = 2
                else:
                    value = digital_value.group()
                title_two_value.append(int(value))

            if result2 and result4:
                title_three_value_group = result4[0]
                title_three_number.append(title_three_value_group)
            if title_three_number != [] and result3:
                title_three_value_group = result3[0]
                digital_value = re.search('\d+', title_three_value_group)
                value = digital_value.group()
                title_three_value.append(int(value))
            if result2 and result3 and not result1:
                if len(result3) == 1:
                    title_two_value_group = result3[0]
                    digital_value = re.search('\d+', title_two_value_group)
                    if digital_value == None:
                        value = -1
                    else:
                        value = digital_value.group()
                    title_two_value.append(int(value))
                else:
                    title_two_number0 = result2[0]
                    digital_title_two_number = re.search('\d+', title_two_number0)
                    digital_title_two_group = digital_title_two_number.group()
                    # title_two_number.append(digital_title_two_group)

                    value = result3[1]
                    digital_value = re.search('\d+', value)
                    if digital_value == None:
                        value0 = -1
                    else:
                        value0 = digital_value.group()

                    # title_two_value.append(int(value0))

                    title_two_number_value_list.append([int(digital_title_two_group), int(value0)])

            if result1 and result2 and result3 and len(title_number) == 1:
                if len(result3) == 1:
                    title_two_value_group = result3[0]
                    digital_value = re.search('\d+', title_two_value_group)
                    if digital_value == None:
                        value = -1
                    else:
                        value = digital_value[0]
                    title_two_value.append(int(value))
                else:
                    title_two_number0 = result2[0]
                    digital_title_two_number = re.search('\d+', title_two_number0)
                    digital_title_two_group = digital_title_two_number.group()
                    # title_two_number.append(digital_title_two_group)

                    value = result3[1]
                    digital_value = re.search('\d+', value)
                    if digital_value == None:
                        value0 = -1
                    else:
                        value0 = digital_value.group()

                    # title_two_value.append(int(value0))

                    title_two_number_value_list.append([int(digital_title_two_group), int(value0)])
            if result5:
                continue
        total_title_quantity = len(title_two_number)
        title_with_value = dict(zip(title_two_number, title_two_value))

        if title_two_number == [] and title_two_value == [] and title_two_number_value_list != []:
            total_title_quantity = len(title_two_number_value_list)
            title_with_value = dict(zip([ele[0] for ele in title_two_number_value_list],
                                        [ele[1] for ele in title_two_number_value_list]))
        elif title_two_number == [] and title_two_value != [] and title_two_number_value_list != []:
            ele_0 = []
            for element in title_two_number_value_list:
                ele_0.append(element[0])
            ele_1 = [ele for ele in range(1, int(len(title_two_number_value_list) + len(title_two_value) + 1))]

            lack_number = list(set(ele_1) - set(ele_0))
            lack_number_with_value = [lack_number[0], title_two_value[0]]
            title_two_number_value_list.append(lack_number_with_value)

            title_two_number_value_list = sorted(title_two_number_value_list, key=lambda k: k[0])

            title_with_value = dict(
                zip([ele[0] for ele in title_two_number_value_list], [ele[1] for ele in title_two_number_value_list]))
            total_title_quantity = len(title_with_value)
        elif title_two_number != [] and title_two_value == [] and title_two_number_value_list == []:
            title_two_value11 = []
            for i in range(len(title_two_number) + 1):
                if i:
                    title_two_value11.append(-1)
            total_title_quantity = len(title_two_number)
            title_with_value = dict(zip(title_two_number, title_two_value11))
        elif title_two_number != [] and title_two_value != [] \
                and title_two_number_value_list == [] and len(title_two_number) == len(title_two_value):
            title_with_value = dict(zip(title_two_number, title_two_value))
            total_title_quantity = len(title_two_number)

        if total_score == [] and title_two_value != []:
            total_score = int(sum(title_two_value))
        elif total_score == [] and title_two_value == []:
            total_score = -1
        else:
            total_score = total_score[0]

        if title_two_value == [] and total_score != -1 and title_two_number != []:
            title_two_value = int(total_score / len(title_two_number))
        elif title_two_value == [] and total_score == -1 and title_two_number != []:
            title_two_value0 = []
            for i in range(len(title_two_number) + 1):
                if i:
                    title_two_value0.append(-1)
            title_with_value = dict(zip(title_two_number, title_two_value0))

        if title_number != [] and total_score != []:
            title_number = title_number[0]
            total_score = total_score
        elif title_number == [] and total_score != []:
            title_number = -1
            total_score = total_score
        elif title_number != [] and total_score == []:
            title_number = title_number[0]
            total_score = -1

        if title_number == -1 and title_two_number == [] and total_score == -1 and title_two_value == [] and total_title_quantity == 0:
            per_title_content = {}
            per_title_content['class_name'] = ele['class_name']
            per_title_content['bounding_box'] = ele['bounding_box']

            per_title_content['title_number'] = -1
            per_title_content['total_title_quantity'] = -1
            per_title_content['total_score'] = -1
            per_title_content['title_with_value'] = -1

            list_of_all.append(per_title_content)
        elif title_number == -1 and title_two_number != [] and total_score == -1 and title_two_value == 0 and total_title_quantity != 0:
            per_title_content = {}

            per_title_content['class_name'] = ele['class_name']
            per_title_content['bounding_box'] = ele['bounding_box']
            per_title_content['title_number'] = -1
            per_title_content['total_title_quantity'] = total_title_quantity
            per_title_content['total_score'] = -1
            per_title_content['title_with_value'] = -1

            list_of_all.append(per_title_content)

        elif total_title_quantity == 0 and title_with_value == []:
            per_title_content = {}

            per_title_content['class_name'] = ele['class_name']
            per_title_content['bounding_box'] = ele['bounding_box']
            per_title_content['title_number'] = title_number
            per_title_content['total_title_quantity'] = -1
            per_title_content['total_score'] = total_score
            per_title_content['title_with_value'] = -1
            list_of_all.append(per_title_content)

        else:
            if total_title_quantity == 0:
                total_title_quantity = -1
            if title_with_value == {}:
                title_with_value = -1
            per_title_content = {}
            per_title_content['class_name'] = ele['class_name']
            per_title_content['bounding_box'] = ele['bounding_box']
            per_title_content['title_number'] = title_number
            per_title_content['total_score'] = total_score
            per_title_content['total_title_quantity'] = total_title_quantity
            per_title_content['title_with_value'] = title_with_value
            list_of_all.append(per_title_content)
    return list_of_all


def get_sheet_points(sheet_dict_list):
    region_dict_front = []
    region_dict_back = []

    for sheet_dict_s in sheet_dict_list:
        regions = sheet_dict_s['sheet_dict']['regions']
        class_name_list = [ele['class_name'] for ele in regions]
        if 'choice_s' in class_name_list or 'choice_m' in class_name_list:
            region_dict_front.append(sheet_dict_s)
        else:
            region_dict_back.append(sheet_dict_s)

    sheet_dict_of_front = []
    sheet_dict_of_front_without_solve_cloze = []
    for sheet_dict_front_s in region_dict_front:
        h, w = sheet_dict_front_s['shape'][0], sheet_dict_front_s['shape'][1]
        words_result_front = sheet_dict_front_s['ocr']
        cloze_or_solve_list_front = []
        if sheet_dict_front_s['sheet_dict']['subject'] != 'english':
            cloze_or_solve_list_front = [ele for ele in sheet_dict_front_s['sheet_dict']['regions']
                                         if ele['class_name'] == 'solve0' or ele['class_name'] == 'solve']
        cloze_s_list = [ele for ele in sheet_dict_front_s['sheet_dict']['regions'] if ele['class_name'] == 'cloze_s']
        cloze_s_box_with_content = []
        for ele in cloze_s_list:
            cloze_s_dict = {}
            cloze_s_region = utils.crop_region(sheet_dict_list[0]['raw_image'], ele['bounding_box'])
            cloze_s_res = get_ocr_text_and_coordinate(cloze_s_region, ocr_accuracy=OCR_ACCURACY, language_type='CHN_ENG')
            if cloze_s_res == []:
                cloze_s_res0 = cloze_s_res
            else:
                cloze_s_res0 = change_box(cloze_s_res, ele['bounding_box'])

            cloze_s_dict['class_name'] = ele['class_name']
            cloze_s_dict['bounding_box'] = ele['bounding_box']
            cloze_s_dict['above_content'] = cloze_s_res0
            cloze_s_box_with_content.append(cloze_s_dict)
        solve_and_cloze_s_list = cloze_or_solve_list_front + cloze_s_list

        if len(solve_and_cloze_s_list) == 0:
            sheet_dict_of_front_without_solve_cloze.append(sheet_dict_front_s)
        else:
            cloze_or_solve0_box_list_front = split_col(solve_and_cloze_s_list, w)

            box_with_content_front = []
            for one_col_box in cloze_or_solve0_box_list_front:

                for ele_box in one_col_box:
                    big_box = {}
                    words_list = []
                    for words_res in words_result_front:
                        xmin = words_res['location']['left']
                        ymin = words_res['location']['top']
                        xmax = words_res['location']['left'] + words_res['location']['width']
                        ymax = words_res['location']['top'] + words_res['location']['height']
                        words_bbox = [xmin, ymin, xmax, ymax]
                        if utils.decide_coordinate_contains1(words_bbox,
                                                             [ele_box['bounding_box']['xmin'],
                                                              ele_box['bounding_box']['ymin'],
                                                              ele_box['bounding_box']['xmax'],
                                                              ele_box['bounding_box']['ymax']]):
                            words_list.append(words_res)
                    big_box['class_name'] = ele_box['class_name']
                    big_box['bounding_box'] = ele_box['bounding_box']
                    big_box['above_content'] = words_list
                    box_with_content_front.append(big_box)
            box_with_content_front1 = box_with_content_front + cloze_s_box_with_content
            list_of_front = get_total_title_quantity_and_value(box_with_content_front1)
            sheet_dict_of_front.append(list_of_front)

    sheet_dict_of_back = []
    sheet_dict_of_back_without_solve_cloze = []
    for single_sheet_dict in region_dict_back:
        h1, w1 = single_sheet_dict['shape'][0], single_sheet_dict['shape'][1]
        words_result_back = single_sheet_dict['ocr']
        cloze_or_solve0_box_list_back = []
        if single_sheet_dict['sheet_dict']['subject'] != 'english':
            cloze_or_solve0_box_list_back = [ele for ele in single_sheet_dict['sheet_dict']['regions']
                                         if ele['class_name'] == 'solve0'
                                         or ele['class_name'] == 'solve']
        cloze_s_list = [ele for ele in single_sheet_dict['sheet_dict']['regions'] if ele['class_name'] == 'cloze_s']
        cloze_s_box_with_content = []
        for ele in cloze_s_list:
            cloze_s_dict = {}
            cloze_s_region = utils.crop_region(sheet_dict_list[0]['raw_image'], ele['bounding_box'])
            cloze_s_res = get_ocr_text_and_coordinate(cloze_s_region, ocr_accuracy=OCR_ACCURACY,
                                                                       language_type='CHN_ENG')
            if cloze_s_res == []:
                cloze_s_res0 = cloze_s_res
            else:
                cloze_s_res0 = change_box(cloze_s_res, ele['bounding_box'])

            cloze_s_dict['class_name'] = ele['class_name']
            cloze_s_dict['bounding_box'] = ele['bounding_box']
            cloze_s_dict['above_content'] = cloze_s_res0
            cloze_s_box_with_content.append(cloze_s_dict)
        solve_and_cloze_s_list_back = cloze_or_solve0_box_list_back + cloze_s_list

        if len(solve_and_cloze_s_list_back) == 0:
            sheet_dict_of_back_without_solve_cloze.append(single_sheet_dict)
        else:
            cloze_or_solve0_box_list_back = split_col(solve_and_cloze_s_list_back, w1)

            box_with_content_back = []
            for one_col_box in cloze_or_solve0_box_list_back:
                for ele_box in one_col_box:
                    words_list = []
                    big_box = {}
                    for words_res in words_result_back:
                        xmin = words_res['location']['left']
                        ymin = words_res['location']['top']
                        xmax = words_res['location']['left'] + words_res['location']['width']
                        ymax = words_res['location']['top'] + words_res['location']['height']
                        words_bbox = [xmin, ymin, xmax, ymax]
                        if utils.decide_coordinate_contains1(words_bbox,
                                                             [ele_box['bounding_box']['xmin'],
                                                              ele_box['bounding_box']['ymin'],
                                                              ele_box['bounding_box']['xmax'],
                                                              ele_box['bounding_box']['ymax']]):
                            words_list.append(words_res)
                    big_box['class_name'] = ele_box['class_name']
                    big_box['bounding_box'] = ele_box['bounding_box']
                    big_box['above_content'] = words_list
                    box_with_content_back.append(big_box)
            # print(box_with_content_back)

            box_with_content_back1 = box_with_content_back + cloze_s_box_with_content
            # box_with_content_back1 = sorted(box_with_content_back1, key=lambda k: k.get('above_content'))
            list_of_back = get_total_title_quantity_and_value(box_with_content_back1)
            sheet_dict_of_back.append(list_of_back)

    if sheet_dict_of_front != [] and sheet_dict_of_front_without_solve_cloze == []:
        sheet_dict_of_front = check_classes(sheet_dict_of_front)
    else:
        sheet_dict_of_front = []
    if sheet_dict_of_back != [] and sheet_dict_of_back_without_solve_cloze == []:
        sheet_dict_of_back = check_classes(sheet_dict_of_back)
    else:
        sheet_dict_of_back = []

    if sheet_dict_of_front != [] and sheet_dict_of_back != []:
        for single_back in sheet_dict_of_back:
            if single_back[0]['title_number'] == -1 and single_back[0]['title_with_value'] != -1:
                if single_back[1]['title_number'] == -1:
                    title_number = single_back[0]['title_number']
                    single_back[0].update({'number': title_number})
                elif single_back[1]['title_number'] != -1:
                    title_number = int(single_back[1]['title_number']) - 1
                    single_back[0].update({'number': title_number})
                else:
                    continue
            else:
                continue

    if sheet_dict_of_front != []:
        for ele11 in region_dict_front:
            class_names_front = ele11['sheet_dict']['regions']
            for single_classes in class_names_front:
                for ele12 in sheet_dict_of_front:
                    for ele13 in ele12:
                        if single_classes['class_name'] == ele13['class_name'] and single_classes['bounding_box'] == \
                                ele13['bounding_box']:
                            title_number = ele13['title_number']
                            total_score = ele13['total_score']
                            total_title_quantity = ele13['total_title_quantity']
                            title_with_value = ele13['title_with_value']
                            single_classes.update({'number': title_number, 'default_points': total_score,
                                                   'total_title_quantity': total_title_quantity,
                                                   'title_with_value': title_with_value})
                        else:
                            continue
    else:
        region_dict_front = sheet_dict_of_front_without_solve_cloze

    if sheet_dict_of_back != []:
        for ele22 in region_dict_back:
            class_names_back = ele22['sheet_dict']['regions']
            for single_class_back in class_names_back:
                for single_paper in sheet_dict_of_back:       #
                    for regions_all in single_paper:
                        if single_class_back['class_name'] == regions_all['class_name'] and \
                                single_class_back['bounding_box'] == regions_all['bounding_box']:
                            title_number0 = regions_all['title_number']
                            total_score0 = regions_all['total_score']
                            total_title_quantity0 = regions_all['total_title_quantity']
                            title_with_value0 = regions_all['title_with_value']
                            single_class_back.update({'number': title_number0,
                                                      'default_points': total_score0,
                                                      'total_title_quantity': total_title_quantity0,
                                                      'title_with_value': title_with_value0})
                        else:
                            continue
    else:
        region_dict_back = sheet_dict_of_back_without_solve_cloze

    tmp = region_dict_front + region_dict_back
    if len(sheet_dict_list) != len(tmp):
        raise ValueError({'答题卡分数识别结果页数丢失: sheet: {}, result: {}'.format(len(sheet_dict_list), len(tmp))})
    else:
        sheet_dict_list = tmp
        # for index, ele in enumerate(sheet_dict_list):
        #     if sheet_dict_list[index]['sheet_dict']['class_name'] == tmp[index]['sheet_dict']['class_name']:
        #         ele['sheet_dict'].update({'regions': tmp[index]['sheet_dict']['regions']})

    return sheet_dict_list
