# coding:utf-8
# brief：根据指定目录下的json配置文件，将其随机分为训练集和测试集，并将路径写到dataSet目录下的txt文件中

import os
import random
import json


# directory-主路径
# fileType-指定文件类型
# fileList-目标类型文件列表（路径+文件名）
def SearchFiles(directory, fileType):
    fileList = []
    for root, subDirs, files in os.walk(directory):
        for fileName in files:
            if fileName.endswith(fileType):
                # json_file = open(directory + '/' + fileName, 'r', encoding='UTF-8')
                # json_data = json.load(json_file)
                # jsonName = json_data['imagePath']
                # jsonName = jsonName.replace('汽车 ', 'car')
                # jsonName = jsonName.replace('(', '_')
                # jsonName = jsonName.replace(')', '')
                # json_data['imagePath'] = jsonName
                # json_file = open(directory + '/' + fileName, 'w', encoding='UTF-8')
                # json.dump(json_data, json_file, indent=2, ensure_ascii=False)
                # print(jsonName)
                fileList.append(fileName)
    # for fileName in fileList:
    #     if fileName.find('汽车 '):
    #         newName = fileName.replace('汽车 ', 'car')
    #         newName = newName.replace('(', '_')
    #         newName = newName.replace(')', '')
    #         os.rename(fileName, newName)
    return fileList


if __name__ == '__main__':
    run_path = 'D:/DeepLearning/pytorch-gpu117/xm_lidar/'
    last_file_path = '/home/zx/doc/private_hub/yolov8/ultralytics-main/examples/train_xm_lidar/labels/'
    # last_file_path = 'D:/DeepLearning/pytorch-gpu117/yolov8_study/train-seg/labels/'
    txt_save_path = run_path + 'dataSet'
    if not os.path.exists(txt_save_path):
        os.makedirs(txt_save_path)

    trainval_percent = 1
    train_percent = 0.9

    total_json = SearchFiles(run_path + 'labels', '.txt')

    print(total_json)

    num = len(total_json)
    list_index = range(num)
    tv = int(num * trainval_percent)
    tr = int(tv * train_percent)
    trainval = random.sample(list_index, tv)
    train = random.sample(trainval, tr)

    file_trainval = open(txt_save_path + '/trainval.txt', 'w')
    file_test = open(txt_save_path + '/test.txt', 'w')
    file_train = open(txt_save_path + '/train.txt', 'w')
    file_val = open(txt_save_path + '/val.txt', 'w')

    for i in list_index:
        name = last_file_path + total_json[i][:-3] + 'jpg\n'
        if i in trainval:
            file_trainval.write(name)
            if i in train:
                file_train.write(name)
            else:
                file_val.write(name)
        else:
            file_test.write(name)

    file_trainval.close()
    file_train.close()
    file_val.close()
    file_test.close()