import os
from PIL import Image
import argparse
import time
import datetime


def main(config):
    src_dir = config.src_dir
    filenames = os.listdir(src_dir)
    trg_dir = config.trg_dir
    if not os.path.exists(trg_dir):
        os.makedirs(trg_dir)

    num_images = len(filenames)

    start_time = time.time()


    for i, filename in enumerate(filenames):
        src_path = os.path.join(src_dir, filename)
        trg_path = os.path.join(trg_dir, filename)
        Image.open(src_path).resize((config.image_size, config.image_size)).save(trg_path)

        if (i+1) % 100 == 0:
            et = time.time() - start_time
            et = str(datetime.timedelta(seconds=et))[:-7]
            print('Elapsed [{}], Processed [{}/{}] '.format(et, i+1, num_images))

    print('Done')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--src_dir', type=str, default='ffhq_gender/female')
    parser.add_argument('--trg_dir', type=str, default='ffhq_gender_256/female')
    parser.add_argument('--image_size', type=int, default=256)

    config = parser.parse_args()
    print(config)
    main(config)

# python resize.py --src_dir ffhq_gender/female --trg_dir ffhq_gender_256/female --image_size 256
# python resize.py --src_dir ffhq_gender/male --trg_dir ffhq_gender_256/male --image_size 256
# zip ffhq_gender_256.zip -r ffhq_gender_256