#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2022-05-18 09:18:02

import random
import argparse
from pathlib import Path

parser = argparse.ArgumentParser(prog='Face dataset Generation')
parser.add_argument('--face_dir', default='/home/jupyter/data/FFHQ/images512x512', type=str,
                                                metavar='PATH', help="Path to save the face images")
# FFHQ: png, Celeba: png
parser.add_argument('--prefix', default='ffhq', type=str, help="Image format of the HR face images")
parser.add_argument('--num_val', default=500, type=int, help="Ratio for Validation set")
parser.add_argument('--seed', default=1234, type=int, help="Random seed")
parser.add_argument('--im_size', default=512, type=int, help="Random seed")
args = parser.parse_args()

base_dir = Path(__file__).resolve().parents[2] / 'files_txt'
if not base_dir.exists():
    base_dir.mkdir()

path_list = sorted([str(x.resolve()) for x in Path(args.face_dir).glob('*.png')])

file_path = base_dir / f"{args.prefix}{args.im_size}.txt"
if file_path.exists():
    file_path.unlink()
with open(file_path, mode='w') as ff:
    for line in path_list: ff.write(line+'\n')

random.seed(args.seed)
random.shuffle(path_list)
num_train = int(len(path_list) - args.num_val)

file_path_train = base_dir / f"{args.prefix}{args.im_size}_train.txt"
if file_path_train.exists():
    file_path_train.unlink()
with open(file_path_train, mode='w') as ff:
    for line in path_list[:num_train]: ff.write(line+'\n')

file_path_val = base_dir / f"{args.prefix}{args.im_size}_val.txt"
if file_path_val.exists():
    file_path_val.unlink()
with open(file_path_val, mode='w') as ff:
    for line in path_list[num_train:]: ff.write(line+'\n')

print('Train / Validation: {:d}/{:d}'.format(num_train, len(path_list)-num_train))