import os
import argparse

str_command = "{0} --startindex {1} --images {2} --samples {3} --folder {4} {5}"

def main():
    
    parser = argparse.ArgumentParser(description="Run and relaunch computation of pbrt-v3")

    parser.add_argument('--pbrt', type=str, help='executable pbrt file', required=True)
    parser.add_argument('--scenes', type=str, help='folder with all pbrt scenes', required=True)
    parser.add_argument('--output', type=str, help='output folder for scenes', required=True)
    parser.add_argument('--samples', type=int, help='number of expected samples for this run', required=True)
    parser.add_argument('--images', type=int, help='number of expected images for this run', required=True)
    parser.add_argument('--file', type=str, help='file with expected pbrt scenes to launch', required=True)

    args = parser.parse_args()

    p_pbrt = args.pbrt
    p_scenes = args.scenes
    p_output = args.output
    p_samples = args.samples
    p_images = args.images
    p_file   = args.file

    with open(p_file, 'r') as f:
        lines = [ l.replace('\n', '') for l in f.readlines() ]

    for l in lines:

        scene_path = os.path.join(p_scenes, l)

        folder = l.split('/')[-1].replace('.pbrt', '')
        print(folder)
        output_folder = os.path.join(p_output, folder)

        # check restart from
        start_index = 0
        current_nb_images = p_images

        if os.path.exists(output_folder):

            n_elements = len(os.listdir(output_folder))

            # update start index and number of expected images now
            start_index = n_elements
            current_nb_images = p_images - n_elements


        if current_nb_images > 0:
            print('Run {0} images with {1} samples using {2} as first index for : {3}'.format(current_nb_images, p_samples, start_index, l))
            main_command = str_command.format(p_pbrt, start_index, current_nb_images, p_samples, p_output, scene_path)
            os.system(main_command)
        else:
            print('All expected images already generated for : {0}'.format(l))

if __name__ == "__main__":
    main()