#!/usr/bin/env python3

import argparse
import os
from subprocess import call


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Generates dataset")
    parser.add_argument('dataset', help='Dataset directory path', type=str)
    parser.add_argument('-d', help='Decals count', type=int, default=8, dest="decals")
    parser.add_argument('-n', help='Noise level', type=float, default=.4, dest="noise")
    args = parser.parse_args()
    
    dataset_path = args.dataset
    ok_path = "ok_" + dataset_path
    damaged_path = "damaged_" + dataset_path
    
    if not os.path.exists(dataset_path):
        raise RuntimeError("fuk")
    
    if not os.path.exists(damaged_path):
        os.makedirs(damaged_path)

    if not os.path.exists(ok_path):
        os.makedirs(ok_path)
    
    dataset_images_paths = [os.path.join(dataset_path, f) for f in os.listdir(dataset_path)]
    ok_images_paths = [os.path.join(ok_path, f) for f in os.listdir(dataset_path)]
    damaged_images_paths = [os.path.join(damaged_path, f) for f in os.listdir(dataset_path)]

    for dataset_image_path, ok_image_path, damaged_image_path in zip(dataset_images_paths, ok_images_paths, damaged_images_paths):
        call(['python', './uss-mariusz.py', dataset_image_path, ok_image_path, damaged_image_path, '-d', str(args.decals), '-n', str(args.noise)])