import os
import glob
import shutil
from os.path import join

import json

import cv2
import cytoolz
from lxml import etree, objectify
import re
import xml.etree.ElementTree as ET

from lxml.objectify import ObjectifiedElement

coco_classes = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light',
    'fire hydrant', 'stop sign', 'parking meter', 'bench',
    'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
    'umbrella', 'handbag', 'tie', 'suitcase',
    'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'wine glass', 'cup', 'fork',
    'knife', 'spoon', 'bowl', 'banana',
    'apple', 'sandwich', 'orange', 'broccoli',
    'carrot', 'hot dog', 'pizza', 'donut',
    'cake', 'chair', 'couch', 'potted plant',
    'bed', 'dining table', 'toilet', 'tv', 'laptop',
    'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator',
    'book', 'clock', 'vase', 'scissors',
    'teddy bear', 'hair drier', 'toothbrush'
]


def convert2xml(root_dir, output_dir=r'CoCoXml'):
    """
    将coco数据集的注释文件由json转换为xml文件，coco数据集目录结构为，root_dir="coco"：
    coco:
        - train2017
            - *.jpg
        - val2017
            - *.jpg
        - test2017[Optional]
            - *.jpg
        - annotations:
            - instances_train2017.json
            - instances_val2017.json

    """

    def instance2xml_base(anno):
        E = objectify.ElementMaker(annotate=False)
        anno_tree = E.annotation(
            E.folder('VOC2014_instance/{}'.format(anno['category_id'])),
            E.filename(anno['file_name']),
            E.source(
                E.database('MS COCO 2014'),
                E.annotation('MS COCO 2014'),
                E.image('Flickr'),
                E.url(anno['coco_url'])
            ),
            E.size(
                E.width(anno['width']),
                E.height(anno['height']),
                E.depth(3)
            ),
            E.segmented(0),
        )
        return anno_tree

    def instance2xml_bbox(anno, bbox_type='xyxy'):
        """bbox_type: xyxy (xmin, ymin, xmax, ymax); xywh (xmin, ymin, width, height)"""
        assert bbox_type in ['xyxy', 'xywh']
        if bbox_type == 'xyxy':
            xmin, ymin, w, h = anno['bbox']
            xmax = xmin + w
            ymax = ymin + h
        else:
            xmin, ymin, xmax, ymax = anno['bbox']
        E = objectify.ElementMaker(annotate=False)
        anno_tree = E.object(
            E.name(anno['category_id']),
            E.bndbox(
                E.xmin(xmin),
                E.ymin(ymin),
                E.xmax(xmax),
                E.ymax(ymax)
            ),
            E.difficult(anno['iscrowd'])
        )
        return anno_tree

    def parse_instance(content, out_dir, output1):
        """
        将coco数据集转换为xml文件

        :param content: coco2017数据集的json格式注释文件加载后的对象
        :param out_dir: 转换结果的输出文件夹， **/train/
        """
        categories = {d['id']: d['name'] for d in content['categories']}
        # merge images and annotations: id in images vs image_id in annotations
        merged_info_list = list(
            map(cytoolz.merge, cytoolz.join('id', content['images'], 'image_id', content['annotations'])))

        # merged_info_list中存储了所有的图片中所有的目标，每一项为一个目标，同一个图片中的多个目标在merged_info_list中对应多个项，以file_name区分图片
        # 此处的cytoolsz.join语句相当于以下功能（但是速度差别极大，可能超过百倍）：
        """
        merged_info_list = []
        for item in content['annotations']:  # 遍历每一项标注信息
            image_id = item['image_id']  # 拿到标注信息中的image_id
            for image in content['images']:  # 遍历每一个图片
                if image['id'] == image_id:  # 如果图片的id等于image_id，则合并图片的信息到标注信息中
                    item.update(image)
                    merged_info_list.append(item)
                    break
        """
        # cytoolsz.join根据content['annotations']中的image_id字段，在content['images']中查找id字段=前述image_id字段的图片，并将找到的图片的项（是一个字典）与
        # convert category id to name
        for instance in merged_info_list:
            instance['category_id'] = categories[instance['category_id']]
        # group by filename to pool all bbox in same file
        for name, groups in cytoolz.groupby('file_name', merged_info_list).items():
            anno_tree: ObjectifiedElement = instance2xml_base(groups[0])  # groups是一个图片中的多个目标信息的列表
            _ = etree.tostring(anno_tree, pretty_print=True).decode('utf-8')
            # if one file have multiple different objects, save it in each category sub-directory
            filenames = []
            # print(groups)
            for group in groups:  # 遍历每一个目标信息
                filenames.append(
                    join(out_dir, re.sub(" ", "_", group['category_id']), os.path.splitext(name)[0] + ".xml"))
                anno_tree.append(instance2xml_bbox(group, bbox_type='xyxy'))
            for filename in filenames:  # 对于不同的目标
                etree.ElementTree(anno_tree).write(filename, pretty_print=True)

            file_no_cate = join(output1, os.path.splitext(name)[0] + ".xml")
            etree.ElementTree(anno_tree).write(file_no_cate, pretty_print=True)
            print("Formating instance xml file {} done!".format(name))

    def convert(json_file, output, output1):
        output_ = join(output_dir, output)
        output1 = join(output_dir, output1)
        if os.path.exists(json_file):
            content = json.load(open(json_file, 'r'))
            sub_dirs = [re.sub(" ", "_", cate['name']) for cate in content['categories']]
            for sub_dir in sub_dirs:
                sub_dir = os.path.join(output_, str(sub_dir))
                if not os.path.exists(sub_dir):
                    os.makedirs(sub_dir)
            if not os.path.exists(output1):
                os.makedirs(output1)
            parse_instance(content, output_, output1)
        else:
            return

    if not os.path.isdir(output_dir):
        output_dir = join(root_dir, output_dir)
    anno_file_train = join(root_dir, "annotations", "instances_train2017.json")
    anno_file_val = join(root_dir, "annotations", "instances_val2017.json")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    convert(anno_file_train, "train_cate", "train_total")
    convert(anno_file_val, "val_cate", "val_total")


def xml2txt(root_dir, classes):
    def convert(size, box):
        # 将图片尺寸和box的坐标信息转为为box在图片中的百分比信息，输出的xywh指box在图片宽度方向的百分比，高度方向的百分比，宽度占图片的百分比和高度占图片的百分比
        dw = 1. / size[0]
        dh = 1. / size[1]
        x = (box[0] + box[2]) / 2.0
        y = (box[1] + box[3]) / 2.0
        w = min(size[0], box[2] - box[0])  # 目标的宽
        h = min(size[1], box[3] - box[1])  # 目标的高
        # print(x,y,w,h)
        x = x * dw
        w = w * dw
        y = y * dh
        h = h * dh
        return (x, y, w, h)

    def convert_annotation(xml, list_file):

        in_file = open(xml, encoding='utf-8')
        tree = ET.parse(in_file)
        root = tree.getroot()

        xmlsize = root.find('size')

        w = int(xmlsize.find('width').text)
        h = int(xmlsize.find('height').text)
        # print(w,h)

        for obj in root.iter('object'):
            difficult = 0
            if obj.find('difficult') != None:
                difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                # print(cls, "------------------------------------\n")
                continue
            cls_id = classes.index(cls)
            # print(cls, cls_id)
            xmlbox = obj.find('bndbox')
            x0 = float(xmlbox.find('xmin').text)
            y0 = float(xmlbox.find('ymin').text)
            x1 = float(xmlbox.find('xmax').text)
            y1 = float(xmlbox.find('ymax').text)
            xmin = min(x0, x1)
            ymin = min(y0, y1)
            xmax = max(x0, x1)
            ymax = max(y0, y1)

            # b = (int(xmin), int(ymin), int(xmax), int(ymax))
            b = (float(xmin), float(ymin), float(xmax), float(ymax))
            if w == 0 or h == 0:
                print('目标宽度为0或高度为0')
            bb = convert((w, h), b)

            list_file.write(" " + str(cls_id) + "," + ",".join([str(a) for a in bb]))

    def inner_xml2txt(coco_xml_path, coco_img_path, mode):
        """
        将xml转换为txt文件，txt文件中，每一行包含多个元素，以空格分开，第一个元素必然是图片路径，后续元素为0个或多个bBoxInfo，其中，bboxinfo
        又是一个以,分割的字符串，包含五个元素，第一个是类别id，后四个是box的坐标信息
        """
        if mode == "train":
            file_write_txt = join(txt_path, 'train2017_v8.txt')
        else:
            file_write_txt = join(txt_path, 'val2017_v8.txt')
        xmls = glob.glob(join(coco_xml_path, '*.xml'))
        list_file = open(file_write_txt, 'w', encoding='utf-8')

        for xml in xmls:
            img = xml.replace(coco_xml_path, coco_img_path).replace('.xml', '.jpg')
            if not os.path.exists(img):
                print(img, ' is not exit')
                continue
            print(img)
            list_file.write(img)
            convert_annotation(xml, list_file)
            list_file.write('\n')
        list_file.close()

    if classes is None:
        classes = coco_classes
    # 图片路径
    coco_img_path_train = join(root_dir, 'train2017')
    coco_img_path_val = join(root_dir, 'val2017')
    # Xml路径
    coco_xml_path_train = join(root_dir, 'CoCoXml', 'train_total')
    coco_xml_path_val = join(root_dir, 'CoCoXml', 'val_total')
    # txt保存路径
    txt_path = join(root_dir, 'CoCoXml', 'labels')
    if not os.path.exists(txt_path):
        os.makedirs(txt_path)

    inner_xml2txt(coco_xml_path_train, coco_img_path_train, "train")
    inner_xml2txt(coco_xml_path_val, coco_img_path_val, "val")


def split_txt(root_dir):
    def split(mode):
        global_txt = join(root_dir, 'CoCoXml', 'labels', f'{mode}2017_v8.txt')
        # 保存txt路径
        save_txt_dir = str(join(root_dir, "yolo", "labels", mode))
        save_img_dir = str(join(root_dir, "yolo", "images", mode))
        if not os.path.exists(save_img_dir):
            os.makedirs(save_img_dir)
        if not os.path.exists(save_txt_dir):
            os.makedirs(save_txt_dir)
        file = open(global_txt, 'r', encoding='utf-8')
        lines = file.readlines()
        for line in lines:
            line = line.split('\n')[0]
            img_file = line.split(' ')[0]
            box_info = line.split(' ')[1:]
            if len(box_info) == 0:
                continue
            # 复制图片至目标文件夹
            img_id = img_file.split('\\')[-1].split(".")[0]
            shutil.copyfile(img_file, join(save_img_dir, img_id + '.jpg'))
            save_txt_file = join(save_txt_dir, f"{img_id}.txt")
            with open(save_txt_file, 'w', encoding='utf-8') as file:
                for i in range(len(box_info)):
                    info = box_info[i].split(',')
                    info1 = ' '.join(info)
                    print(info1)
                    file.write(info1 + '\n')

    split("train")
    split("val")


def show_pic_with_label(root_dir, classes=None):
    """
    显示图片和标注，用以验证标注是否正确
    """
    if classes is None:
        classes = coco_classes
    img_dir = join(root_dir, 'yolo', 'images', 'train')
    txt_dir = join(root_dir, 'yolo', 'labels', 'train')
    list1 = glob.glob(os.path.join(img_dir, '*.jpg'))
    list2 = glob.glob(os.path.join(txt_dir, '*.txt'))
    idx = 0
    # print(len(list1))
    # print(len(list2))
    while idx < len(list1):
        imgpath = list1[idx]
        txtpath = list2[idx]
        print(imgpath)
        print(txtpath)
        f = open(txtpath, 'r')
        lines = f.readlines()
        img = cv2.imread(imgpath)
        h, w, c = img.shape
        colors = [
                     [0, 0, 255], [0, 255, 0], [0, 255, 255], [255, 0, 0], [255, 255, 0],
                     [0, 0, 128], [0, 128, 0], [0, 128, 128], [128, 0, 0], [128, 128, 0]
                 ] * 8

        colors = colors[:len(classes)]  # 有多少个类，则使用多少个颜色，最多不超过80个类

        # print(lines)
        for line in lines:
            # print(line)
            l = line.split(' ')
            # print(len(l))
            label = l[0]
            cx = float(l[1]) * w
            cy = float(l[2]) * h
            weight = float(l[3]) * w
            height = float(l[4]) * h
            xmin = cx - weight / 2
            ymin = cy - height / 2
            xmax = cx + weight / 2
            ymax = cy + height / 2
            print(label, (xmin, ymin), (xmax, ymax))
            color = colors[int(label)]
            cv2.putText(img, label, (int(xmin), int(ymin)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 1, cv2.LINE_AA)
            cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)
            kpts = []

        img = cv2.resize(img, None, fx=1, fy=1)
        cv2.imshow('1', img)
        key = cv2.waitKey(0)
        if key == ord('q'):
            break
        if key == ord('z'):
            idx -= 1
        else:
            idx += 1


def transfer_coco_to_yolo(root_dir, classes=None):
    """
    classes可以设置需要训练的模型识别的目标种类，则只生成相应目标的训练数据集
    """
    convert2xml(root_dir)
    xml2txt(root_dir, classes)
    split_txt(root_dir)
    # show_pic_with_label(root_dir, classes)


if __name__ == "__main__":
    root = r"D:\dataset"
    # transfer_coco_to_yolo(root)
    show_pic_with_label(root)
