import os
import json
import argparse
import numpy as np
import glob
import cv2
from sklearn.model_selection import train_test_split
from labelme import utils
from tqdm import tqdm

np.random.seed(41)

# 0为背景
classname_to_id = {
    "front": 1, # 从1开始标注
    "back": 2,
}


class Lableme2CoCo:

    def __init__(self):
        self.images = []
        self.annotations = []
        self.categories = []
        self.img_id = 0
        self.ann_id = 0

    def save_coco_json(self, instance, save_path):
        json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)  # indent=2 更加美观显示

    # 由json文件构建COCO
    def to_coco(self, json_path_list):
        self._init_categories()
        for json_path in json_path_list:
            obj = self.read_jsonfile(json_path)
            self.images.append(self._image(obj, json_path))
            shapes = obj['shapes']
            for shape in shapes:
                annotation = self._annotation(shape)
                self.annotations.append(annotation)
                self.ann_id += 1
            self.img_id += 1
        instance = {}
        instance['info'] = 'spytensor created'
        instance['license'] = ['license']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    # 构建类别
    def _init_categories(self):
        for k, v in classname_to_id.items():
            category = {}
            category['id'] = v
            category['name'] = k
            self.categories.append(category)

    # 构建COCO的image字段
    def _image(self, obj, path):
        image = {}
        # img_x = utils.img_b64_to_arr(obj['imageData'])
        # h, w = img_x.shape[:-1]
        image['height'] = obj['imageHeight']
        image['width'] = obj['imageWidth']
        # image['height'] = h
        # image['width'] = w
        image['id'] = self.img_id
        image['file_name'] = os.path.basename(path).replace(".json", ".png")
        return image

    # 构建COCO的annotation字段
    def _annotation(self, shape):
        # print('shape', shape)
        label = shape['label']
        points = shape['points']
        annotation = {}
        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        annotation['category_id'] = int(classname_to_id[label])
        annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
        annotation['bbox'] = self._get_box(points)
        annotation['iscrowd'] = 0
        annotation['area'] = 1.0
        return annotation

    # 读取json文件，返回一个json对象
    def read_jsonfile(self, path):
        with open(path, "r", encoding='utf-8') as f:
            return json.load(f)

    # COCO的格式： [x1,y1,w,h] 对应COCO的bbox格式
    def _get_box(self, points):
        min_x = min_y = np.inf
        max_x = max_y = 0
        for x, y in points:
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
        return [min_x, min_y, max_x - min_x, max_y - min_y]


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--anno_dirs', type=str, nargs="+", default=['./'])
    parser.add_argument('--train_ratio', type=float, default=0.9)
    args = parser.parse_args()
    labelme_folds = args.anno_dirs
    json_list_path = []
    train_path = []
    val_path = []
    # 每个文件夹按照比例划分，遍历完，最后合并
    for labelme_path in labelme_folds:
        list_path = glob.glob(labelme_path + "/*.json")
        json_list_path.extend(list_path)
        train_path1, val_path1 = train_test_split(list_path, test_size=1-args.train_ratio, train_size=args.train_ratio)
        train_path.extend(train_path1)
        val_path.extend(val_path1)
    saved_coco_path = "./"
    print('reading...')
    # 创建文件
    if not os.path.exists("%scoco/train/" % saved_coco_path):
        os.makedirs("%scoco/train/" % saved_coco_path)
    if not os.path.exists("%scoco/train/images/" % saved_coco_path):
        os.makedirs("%scoco/train/images/" % saved_coco_path)
    if not os.path.exists("%scoco/eval/" % saved_coco_path):
        os.makedirs("%scoco/eval/" % saved_coco_path)
    if not os.path.exists("%scoco/eval/images/" % saved_coco_path):
        os.makedirs("%scoco/eval/images/" % saved_coco_path)
    print('total images: ', len(json_list_path))
    # 数据划分,这里没有区分val2017和tran2017目录，所有图片都放在images目录下

    print("train_n:", len(train_path), 'val_n:', len(val_path))

    # 把训练集转化为COCO的json格式
    l2c_train = Lableme2CoCo()
    train_instance = l2c_train.to_coco(train_path)
    l2c_train.save_coco_json(train_instance, '%scoco/train/annotations.json' % saved_coco_path)
    print("train images: %d" % len(train_path))
    for file in tqdm(train_path):
        img_name = file.replace('json', 'png')
        temp_img = cv2.imread(img_name)
        try:
            img_name = str(img_name).split('\\')[-1]
            cv2.imwrite("{}coco/train/images/{}".format(saved_coco_path, img_name.replace('jpg', 'jpg')), temp_img)
        except Exception as e:
            print(e)
            print('Wrong Image:', img_name)
            continue
        # print(img_name + '-->', img_name.replace('jpg', 'jpg'))
    print("eval images: %d" % len(val_path))
    for file in tqdm(val_path):
        img_name = file.replace('json', 'png')
        temp_img = cv2.imread(img_name)
        try:
            img_name = str(img_name).split('\\')[-1]
            cv2.imwrite("{}coco/eval/images/{}".format(saved_coco_path, img_name.replace('jpg', 'jpg')), temp_img)
        except Exception as e:
            print(e)
            print('Wrong Image:', img_name)
            continue

    # 把验证集转化为COCO的json格式
    l2c_val = Lableme2CoCo()
    val_instance = l2c_val.to_coco(val_path)
    l2c_val.save_coco_json(val_instance, '%scoco/eval/annotations.json' % saved_coco_path)
    with open('./coco/about.txt', 'w') as f:
        f.write(str(args.anno_dirs))
    f.close()