import copy
import json
import os
import random
import uuid

import cv2
import re
import pandas as pd
import numpy as np
import jieba
from shapely.geometry import Polygon, MultiPoint
from tools import get_file_paths, load_json
from word2vec import jwq_word2vec, simple_word2vec, jieba_and_tencent_word2vec

def bbox_iou(go_bbox, label_bbox, mode='iou'):
    # 所有点的最小凸的表示形式，四边形对象，会自动计算四个点，最后顺序为：左上 左下  右下 右上 左上
    go_poly = Polygon(go_bbox).convex_hull
    label_poly = Polygon(label_bbox).convex_hull
    if not go_poly.is_valid or not label_poly.is_valid:
        print('formatting errors for boxes!!!! ')
        return 0
    if go_poly.area == 0 or label_poly.area == 0 :
        return 0

    inter = Polygon(go_poly).intersection(Polygon(label_poly)).area
    go_area = Polygon(go_poly).area

    return inter / go_area
    
    # if mode == 'iou':
    #     union = go_poly.area + label_poly.area - inter
    # elif mode =='tiou':
    #     union_poly = np.concatenate((go_bbox, label_bbox))   #合并两个box坐标，变为8*2
    #     union = MultiPoint(union_poly).convex_hull.area
    #     # coors = MultiPoint(union_poly).convex_hull.wkt
    # elif mode == 'giou':
    #     union_poly = np.concatenate((go_bbox, label_bbox))
    #     union = MultiPoint(union_poly).envelope.area
    #     # coors = MultiPoint(union_poly).envelope.wkt
    # elif mode == 'r_giou':
    #     union_poly = np.concatenate((go_bbox, label_bbox))
    #     union = MultiPoint(union_poly).minimum_rotated_rectangle.area
    #     # coors = MultiPoint(union_poly).minimum_rotated_rectangle.wkt
    # else:
    #     raise Exception('incorrect mode!')

    # if union == 0:
    #     return 0
    # else:
    #     return inter / union

def clean_go_res(go_res_dir):
    go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
    for go_res_json_path in go_res_json_paths:
        print('Info: start {0}'.format(go_res_json_path))
    
        remove_idx_set = set()
        src_go_res_list = load_json(go_res_json_path)
        for idx, (_, text) in enumerate(src_go_res_list):
            if text.strip() == '':
                remove_idx_set.add(idx)
                print(text)
        
        if len(remove_idx_set) > 0:
            for del_idx in remove_idx_set:
                del src_go_res_list[del_idx]
        
        go_res_list = sorted(src_go_res_list, key=lambda x: (x[0][1], x[0][0]), reverse=False)

        with open(go_res_json_path, 'w') as fp:
            json.dump(go_res_list, fp)
            print('Rerewirte {0}'.format(go_res_json_path))

def char_length_statistics(go_res_dir):
    max_char_length = None
    target_file_name = None
    go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
    for go_res_json_path in go_res_json_paths:
        print('Info: start {0}'.format(go_res_json_path))
        src_go_res_list = load_json(go_res_json_path)
        for _, text in src_go_res_list:
            if max_char_length is None or len(text.strip()) > max_char_length:
                max_char_length = len(text.strip())
                target_file_name = go_res_json_path
    return max_char_length, target_file_name

def char_length_statistics_jieba(go_res_dir):
    max_char_length = None
    target_file_name = None
    max_char_length = None
    statistics_dict = {}
    go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
    for go_res_json_path in go_res_json_paths:
        print('Info: start {0}'.format(go_res_json_path))
        src_go_res_list = load_json(go_res_json_path)
        for _, text in src_go_res_list:
            jieba_char_list = list(filter(lambda x:re.match(r'[\u4e00-\u9fa5]', x), jieba.lcut(text.strip())))
            length = len(jieba_char_list) 
            if length in statistics_dict:
                statistics_dict[length] += 1
            else:
                statistics_dict[length] = 1
            if max_char_length is None or length > max_char_length:
                target_file_name = go_res_json_path
                target_jieba_char_list = jieba_char_list 
                max_char_length = length
    return max_char_length, target_file_name, target_jieba_char_list, statistics_dict 

def bbox_statistics(go_res_dir):
    max_seq_count = None
    seq_sum = 0
    file_count = 0

    go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
    for go_res_json_path in go_res_json_paths:
        print('Info: start {0}'.format(go_res_json_path)) 

        go_res_list = load_json(go_res_json_path)
        seq_sum += len(go_res_list)
        file_count += 1
        if max_seq_count is None or len(go_res_list) > max_seq_count:
            max_seq_count = len(go_res_list)
            max_seq_file_name = go_res_json_path

    seq_lens_mean = seq_sum // file_count 
    return max_seq_count, seq_lens_mean, max_seq_file_name

def text_statistics(go_res_dir):
    """
    Args:
        go_res_dir: str 通用OCR的JSON文件夹
    Returns: list 出现次数最多的文本及其次数
    """
    json_count = 0
    text_dict = {}
    go_res_json_paths = get_file_paths(go_res_dir, ['.json', ])
    for go_res_json_path in go_res_json_paths:
        print('Info: start {0}'.format(go_res_json_path))
        json_count += 1
        go_res = load_json(go_res_json_path)
        for _, text in go_res.values():
            if text in text_dict:
                text_dict[text] += 1
            else:
                text_dict[text] = 1
    top_text_list = []
    # 按照次数排序
    for text, count in sorted(text_dict.items(), key=lambda x: x[1], reverse=True):
        if text == '':
            continue
        # 丢弃：次数少于总数的2/3
        if count <= json_count // 3:
            break
        top_text_list.append((text, count))
    return top_text_list

def build_anno_file(dataset_dir, anno_file_path):
    img_list = os.listdir(dataset_dir)
    random.shuffle(img_list)
    df = pd.DataFrame(columns=['name'])
    df['name'] = img_list
    df.to_csv(anno_file_path)

def build_dataset(img_dir, go_res_dir, label_dir, top_text_list, skip_list, save_dir, is_create_map=False):
    """
    Args:
        img_dir: str 图片目录
        go_res_dir: str 通用OCR的JSON保存目录
        label_dir: str 标注的JSON保存目录
        top_text_list: list 出现次数最多的文本及其次数
        skip_list: list 跳过的图片列表
        save_dir: str 数据集保存目录
    """
    if os.path.exists(save_dir):
        return
    else:
        os.makedirs(save_dir, exist_ok=True)

    # 开票日期 发票代码 机打号码 车辆类型 电话 
    # 发动机号码 车架号 帐号 开户银行 小写
    group_cn_list = ['开票日期', '发票代码', '机打号码', '车辆类型', '电话', '发动机号码', '车架号', '帐号', '开户银行', '小写']
    test_group_id = [1, 2, 5, 9, 20, 15, 16, 22, 24, 28]

    create_map = {}
    for img_name in sorted(os.listdir(img_dir)):
        if img_name in skip_list:
            print('Info: skip {0}'.format(img_name))
            continue

        print('Info: start {0}'.format(img_name))
        image_path = os.path.join(img_dir, img_name)
        img = cv2.imread(image_path)
        h, w, _ = img.shape
        base_image_name, _ = os.path.splitext(img_name)
        go_res_json_path = os.path.join(go_res_dir, '{0}.json'.format(base_image_name))
        go_res_list = load_json(go_res_json_path)

        valid_lens = len(go_res_list)

        top_text_idx_set = set()
        for top_text, _ in top_text_list:
            for go_idx, (_, text) in enumerate(go_res_list):
                if text == top_text:
                    top_text_idx_set.add(go_idx)
                    break

        label_json_path = os.path.join(label_dir, '{0}.json'.format(base_image_name))
        label_res = load_json(label_json_path)

        group_list = []
        for group_id in test_group_id:
            for item in label_res.get("shapes", []):
                if item.get("group_id") == group_id:
                    label_bbox = list()
                    for point in item['points']:
                        label_bbox.extend(point)
                    group_list.append(label_bbox)
                    break
            else:
                group_list.append(None)
        
        label_idx_dict = dict()
        for label_idx, label_bbox in enumerate(group_list):
            if isinstance(label_bbox, list):
                for go_idx, (go_bbox, _) in enumerate(go_res_list):
                    if go_idx in top_text_idx_set or go_idx in label_idx_dict:
                        continue
                    go_bbox_rebuild = [
                        [go_bbox[0], go_bbox[1]],
                        [go_bbox[2], go_bbox[3]],
                        [go_bbox[4], go_bbox[5]],
                        [go_bbox[6], go_bbox[7]],
                    ]
                    label_bbox_rebuild = [
                        [label_bbox[0], label_bbox[1]],
                        [label_bbox[2], label_bbox[1]],
                        [label_bbox[2], label_bbox[3]],
                        [label_bbox[0], label_bbox[3]],
                    ]
                    iou = bbox_iou(go_bbox_rebuild, label_bbox_rebuild)
                    if iou >= 0.2:
                        label_idx_dict[go_idx] = label_idx 
        
        X = list()
        y_true = list()

        X_no_text = list()

        # dim = 1 + 5 + 8

        # text_vec_max_lens = 15 * 50
        # dim = 1 + 5 + 8 + text_vec_max_lens 

        max_jieba_char = 4
        text_vec_max_lens = max_jieba_char * 100
        dim = 1 + 5 + 8 + text_vec_max_lens 

        num_classes = 10
        for i in range(160):
            if i >= valid_lens:
                X.append([0. for _ in range(dim)])
                y_true.append([0 for _ in range(num_classes)])

                X_no_text.append([0. for _ in range(dim)])

            elif i in top_text_idx_set:
                (x0, y0, x1, y1, x2, y2, x3, y3), text = go_res_list[i]
                feature_vec = [1.]
                feature_vec.extend(simple_word2vec(text))
                feature_vec.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
                # feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
                feature_vec.extend(jieba_and_tencent_word2vec(text, max_jieba_char))
                X.append(feature_vec)

                y_true.append([0 for _ in range(num_classes)])

                feature_vec_no_text = [1.]
                feature_vec_no_text.extend([0. for _ in range(5)])
                feature_vec_no_text.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
                # feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
                feature_vec_no_text.extend([0. for _ in range(text_vec_max_lens)])
                X_no_text.append(feature_vec_no_text)                

            elif i in label_idx_dict:
                (x0, y0, x1, y1, x2, y2, x3, y3), text = go_res_list[i]
                feature_vec = [0.]
                feature_vec.extend(simple_word2vec(text))
                feature_vec.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
                # feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
                feature_vec.extend(jieba_and_tencent_word2vec(text, max_jieba_char))
                X.append(feature_vec)

                base_label_list = [0 for _ in range(num_classes)]
                base_label_list[label_idx_dict[i]] = 1
                y_true.append(base_label_list)

                feature_vec_no_text = [0.]
                feature_vec_no_text.extend([0. for _ in range(5)])
                feature_vec_no_text.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
                # feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
                feature_vec_no_text.extend([0. for _ in range(text_vec_max_lens)])
                X_no_text.append(feature_vec_no_text) 

            else:
                (x0, y0, x1, y1, x2, y2, x3, y3), text = go_res_list[i]
                feature_vec = [0.]
                feature_vec.extend(simple_word2vec(text))
                feature_vec.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
                # feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
                feature_vec.extend(jieba_and_tencent_word2vec(text, max_jieba_char))
                X.append(feature_vec)

                y_true.append([0 for _ in range(num_classes)])

                feature_vec_no_text = [0.]
                feature_vec_no_text.extend([0. for _ in range(5)])
                feature_vec_no_text.extend([(x0/w), (y0/h), (x1/w), (y1/h), (x2/w), (y2/h), (x3/w), (y3/h)])
                # feature_vec.extend(jwq_word2vec(text, text_vec_max_lens))
                feature_vec_no_text.extend([0. for _ in range(text_vec_max_lens)])
                X_no_text.append(feature_vec_no_text) 

        all_data = [X, y_true, valid_lens]
        all_data_no_text = [X_no_text, y_true, valid_lens]

        save_json_name = '{0}.json'.format(uuid.uuid3(uuid.NAMESPACE_DNS, img_name))
        with open(os.path.join(save_dir, save_json_name), 'w') as fp:
            json.dump(all_data, fp)

        save_json_name_2 = '{0}.json'.format(uuid.uuid3(uuid.NAMESPACE_DNS, '{0}_no_text'.format(img_name)))
        with open(os.path.join(save_dir, save_json_name_2), 'w') as fp:
            json.dump(all_data_no_text, fp)

        if is_create_map:
            create_map[img_name] = {
                'x_y_valid_lens': save_json_name, 
                'find_top_text': [go_res_list[i][-1] for i in top_text_idx_set],
                'find_value': {go_res_list[k][-1]: group_cn_list[v] for k, v in label_idx_dict.items()}
            }

        # break
    
    # print(create_map)
    # print(is_create_map)
    if create_map:
        # print(create_map)
        with open(os.path.join(os.path.dirname(save_dir), 'create_map.json'), 'w') as fp:
            json.dump(create_map, fp) 

        # print('top text find:')
        # for i in top_text_idx_set:
        #     _, text = go_res_list[i]
        #     print(text)

        # print('-------------')
        # print('label value find:')
        # for k, v in label_idx_dict.items():
        #     _, text = go_res_list[k]
        #     print('{0}: {1}'.format(group_cn_list[v], text)) 

        # break


if __name__ == '__main__':
    base_dir = '/Users/zhouweiqi/Downloads/gcfp/data'
    go_dir = os.path.join(base_dir, 'go_res')
    dataset_save_dir = os.path.join(base_dir, 'dataset160x414x10-no-text')
    label_dir = os.path.join(base_dir, 'labeled')

    train_go_path = os.path.join(go_dir, 'train')
    train_image_path = os.path.join(label_dir, 'train', 'image')
    train_label_path = os.path.join(label_dir, 'train', 'label')
    train_dataset_dir = os.path.join(dataset_save_dir, 'train')
    train_anno_file_path = os.path.join(dataset_save_dir, 'train.csv')

    valid_go_path = os.path.join(go_dir, 'valid')
    valid_image_path = os.path.join(label_dir, 'valid', 'image')
    valid_label_path = os.path.join(label_dir, 'valid', 'label')
    valid_dataset_dir = os.path.join(dataset_save_dir, 'valid')
    valid_anno_file_path = os.path.join(dataset_save_dir, 'valid.csv')

    # max_seq_lens, seq_lens_mean, max_seq_file_name = bbox_statistics(go_dir)
    # print(max_seq_lens) # 152
    # print(max_seq_file_name) # train/CH-B101805176_page_2_img_0.json
    # print(seq_lens_mean) # 92

    # max_char_lens, target_file_name = char_length_statistics(go_dir)
    # print(max_char_lens) # 72
    # print(target_file_name) # train/CH-B103053828-4.json

    # max_char_length, target_file_name, target_jieba_char_list, statistics_dict = char_length_statistics_jieba(go_dir)
    # print(max_char_length) # 24
    # print(target_file_name) # train/CH-B102551568-6.json
    # print(target_jieba_char_list)
    # print(statistics_dict) # {2: 12077, 1: 12751, 0: 13073, 3: 4423, 4: 1212, 5: 969, 6: 744, 7: 524, 8: 199, 10: 45, 12: 9, 18: 44, 9: 109, 11: 19, 13: 4, 16: 4, 21: 2, 19: 2, 15: 8, 17: 7, 14: 3, 20: 1, 24: 1}

    # top_text_list = text_statistics(go_dir)
    # for t in top_text_list:
    #     print(t)

    filter_from_top_text_list = [
        ('机器编号', 496),
        ('购买方名称', 496),
        ('合格证号', 495),
        ('进口证明书号', 495),
        ('机打代码', 494),
        ('车辆类型', 492),
        ('完税凭证号码', 492),
        ('机打号码', 491),
        ('发动机号码', 491),
        ('主管税务', 491),
        ('价税合计', 489),
        ('机关及代码', 489),
        ('销货单位名称', 486),
        ('厂牌型号', 485),
        ('产地', 485),
        ('商检单号', 483),
        ('电话', 476),
        ('开户银行', 472),
        ('车辆识别代号/车架号码', 463),
        ('身份证号码', 454),
        ('吨位', 452),
        ('备注:一车一票', 439),
        ('地', 432),
        ('账号', 431),
        ('统一社会信用代码/', 424),
        ('限乘人数', 404),
        ('税额', 465),
        ('址', 392)
    ]

    skip_list_train = [
        'CH-B101910792-page-12.jpg',
        'CH-B101655312-page-13.jpg',
        'CH-B102278656.jpg',
        'CH-B101846620_page_1_img_0.jpg',
        'CH-B103062528-0.jpg',
        'CH-B102613120-3.jpg',
        'CH-B102997980-3.jpg',
        'CH-B102680060-3.jpg',
        # # 'CH-B102995500-2.jpg',  # 没value
    ]

    skip_list_valid = [
        # 'CH-B102897920-2.jpg',
        # 'CH-B102551284-0.jpg',
        # 'CH-B102879376-2.jpg',
        # 'CH-B101509488-page-16.jpg',
        # 'CH-B102708352-2.jpg',
    ]

    build_dataset(train_image_path, train_go_path, train_label_path, filter_from_top_text_list, skip_list_train, train_dataset_dir)
    build_anno_file(train_dataset_dir, train_anno_file_path)

    build_dataset(valid_image_path, valid_go_path, valid_label_path, filter_from_top_text_list, skip_list_valid, valid_dataset_dir, True)
    build_anno_file(valid_dataset_dir, valid_anno_file_path)

    # print(simple_word2vec(' fd2jk接口 额24；叁‘，。测ADF壹试！¥？ '))
    # print(jwq_word2vec('发', 15*50))