import glob
import json
import shutil
from os.path import join, dirname, abspath
import os


def xyxy2xc_y_cwh(x1, y1, x2, y2):
    """
    将xyxy形式的标记框转换为(xcenter,ycenter,width,height)形式的标记框
    """
    xc = int((x1 + x2) / 2)
    yc = int((y1 + y2) / 2)
    width = abs(x2 - x1)
    height = abs(y2 - y1)
    return xc, yc, width, height


def scale1(bbox, w, h):
    """
    将图片标注归一化
    """
    x = float(bbox[0] / w)
    y = float(bbox[1] / h)
    w = float(bbox[2] / w)
    h = float(bbox[3] / h)
    return x, y, w, h


def load_label(root_dir=r'C:\Users\54067\Pictures\wow3截图', classes=None, splitter=[0.8, 0.2]):
    """
    yolo v8格式的数据是(x,y,w,h)的归一化标记框，其中，(x,y)是标记框的中心坐标
    """
    if classes is None:
        classes = []
    classes_dict = {v: k for k, v in enumerate(classes)}

    train_image_dir = join(root_dir, 'yolo', 'images', 'train')
    val_image_dir = join(root_dir, 'yolo', 'images', 'val')
    train_label_dir = join(root_dir, 'yolo', 'labels', 'train')
    val_label_dir = join(root_dir, 'yolo', 'labels', 'val')
    if not os.path.exists(train_image_dir):
        os.makedirs(train_image_dir)
    if not os.path.exists(val_image_dir):
        os.makedirs(val_image_dir)
    if not os.path.exists(train_label_dir):
        os.makedirs(train_label_dir)
    if not os.path.exists(val_label_dir):
        os.makedirs(val_label_dir)

    total_num = len(glob.glob(join(root_dir, '*.json')))
    train_num = int(total_num * splitter[0])

    i = 0
    for file in glob.glob(join(root_dir, '*.json')):
        basename = os.path.splitext(os.path.basename(file))[0]
        source_image_file = join(root_dir, f"{basename}.jpg")
        if i < train_num:
            out_image_file = join(train_image_dir, f"{basename}.jpg")
            out_txt_file = join(train_label_dir, f"{basename}.txt")
        else:
            out_image_file = join(val_image_dir, f"{basename}.jpg")
            out_txt_file = join(val_label_dir, f"{basename}.txt")
        i += 1
        label = json.load(open(file, 'r', encoding='utf-8'))
        width_pic = label['imageWidth']
        height_pic = label['imageHeight']
        with open(out_txt_file, 'w', encoding='utf-8') as f:
            for shape in label['shapes']:
                shape: dict
                points = shape['points']
                x1, y1 = points[0]
                x2, y2 = points[1]
                _ = xyxy2xc_y_cwh(x1, y1, x2, y2)
                bbox = scale1(_, width_pic, height_pic)
                cate_id = classes_dict[shape['label']]
                f.write(f"{cate_id} {bbox[0]} {bbox[1]} {bbox[2]} {bbox[3]}\n")
        shutil.copyfile(source_image_file, out_image_file)
    return join(root_dir, 'yolo')


if __name__ == '__main__':
    classes = ["沙漠鞭尾蝎", "怪物1", "狮子宠物"]
    load_label(classes=classes)
