import os 
import glob 
import shutil 
import argparse
import random
from pathlib import Path

random.seed(1000)

def copy_data(src_paths, tgt_dir):
    for p in src_paths:
        txt_path = str(Path(p).with_suffix(".txt"))
        shutil.copy(txt_path, tgt_dir)
        shutil.copy(p, tgt_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--path", type=str, required=True)
    parser.add_argument("--out", type=str, required=True)
    parser.add_argument("--test_ratio", type=float, default=0.1)
    args = parser.parse_args()
    img_paths = [p for p in sorted(glob.glob(args.path + "/*")) if ".txt" not in p]
    random.shuffle(img_paths)
    total_imgs = len(img_paths)
    num_test = int(args.test_ratio * total_imgs)
    if num_test == 0:
        test_paths = []
        train_paths = img_paths
    else:
        train_paths = img_paths[:-num_test]
        test_paths = img_paths[-num_test:]

    out_dir = Path(args.out)
    out_train = out_dir / 'train'
    out_test = out_dir / 'test'
    if not out_train.exists():
        out_train.mkdir(parents=True)
        out_test.mkdir(parents=True)

    copy_data(train_paths, out_train)
    copy_data(test_paths, out_test)
    
    

