import os
import random
import sys
import argparse

parser = argparse.ArgumentParser(description='Split VOC Dataset')
parser.add_argument('-r', '--voc_root', type=str, default='D:\\ai_data\\HardHatWorker_voc\\VOCdevkit\\VOC2007',
                    help='voc root eg:xxx/xxx/VOCdevkit/VOC2007(2012)')
parser.add_argument('-tvp', '--train_val_percent', type=float, default=0.9,
                    help='train val percent')
parser.add_argument('-tp', '--train_percent', type=float, default=0.8,
                    help='train percent')
args = parser.parse_args()
if args.voc_root == "":
    print('VOC Path Is Empty!')
    exit(0)
print(f'VOC Root Path:{args.voc_root}')

if sys.version_info[0] == 2:
    import xml.etree.cElementTree as ET
else:
    import xml.etree.ElementTree as ET


def generate_train_val_test_txt(root: str,train_percent: float,train_val_percent: float):

    assert os.path.exists(root) is True, 'root path is error'

    xml_path = os.path.join(root, 'Annotations')
    assert os.path.exists(xml_path) is True, 'xml path is error'

    image_set_path = os.path.join(root, 'ImageSets')
    if os.path.exists(image_set_path) is False:
        os.mkdir(image_set_path)

    save_path = os.path.join(image_set_path, 'Main')
    if os.path.exists(save_path) is False:
        os.mkdir(save_path)

    total_xml = os.listdir(xml_path)
    num = len(total_xml)
    list_ = range(num)

    train_val_num = int(num * train_val_percent)
    train_num = int(train_val_num * train_percent)

    train_val = random.sample(list_, train_val_num)
    train = random.sample(train_val, train_num)

    print(f'train and val size:{len(train_val)}')
    print(f'train size:{len(train)}')

    file_train_val = open(os.path.join(save_path, 'trainval.txt'), 'w')
    file_test = open(os.path.join(save_path, 'test.txt'), 'w')
    file_train = open(os.path.join(save_path, 'train.txt'), 'w')
    file_val = open(os.path.join(save_path, 'val.txt'), 'w')

    for i in list_:
        xml_name = total_xml[i][:-4]
        if i in train_val:
            file_train_val.write(xml_name + '\n')
            if i in train:
                file_train.write(xml_name + '\n')
            else:
                file_val.write(xml_name + '\n')
        else:
            file_test.write(xml_name + '\n')
    file_train_val.close()
    file_train.close()
    file_test.close()
    file_val.close()
    print('Finished!')


def main(args):
    generate_train_val_test_txt(root=args.voc_root,
                                train_percent=args.train_percent,
                                train_val_percent=args.train_val_percent)


if __name__ == '__main__':
    main(args)
