import os
from glob import glob
import sys
import shutil

from sklearn.model_selection import train_test_split

if __name__ == "__main__":
    input_image_dir = sys.argv[1]

    train_dir = os.path.join(input_image_dir, "train")
    test_dir = os.path.join(input_image_dir, "test")

    if os.path.isdir(train_dir):
        shutil.rmtree(train_dir)

    if os.path.isdir(test_dir):
        shutil.rmtree(test_dir)

    os.makedirs(train_dir)
    os.makedirs(test_dir)

    image_glob = sorted(glob(os.path.join(input_image_dir, "*.jpg")))
    label_glob = sorted(glob(os.path.join(input_image_dir, "*.txt")))

    train_img, test_img, train_label, test_label = train_test_split(
        image_glob, label_glob, train_size=0.9, shuffle=True
    )

    for img, label in zip(train_img, train_label):
        shutil.copy(img, train_dir)
        shutil.copy(label, train_dir)

    for img, label in zip(test_img, test_label):
        shutil.copy(img, test_dir)
        shutil.copy(label, test_dir)

    

