import os
import glob
import numpy as np

def random_choice_file(dir_):
    image_f = np.random.choice(image_lst, size=1, replace=False)[0]


def f():
    # 两个变量，选中的概率分别为2/3, 1/3。选中之后，分别打上tag
    image_index = np.random.choice([0, 1], size=1, replace=False, p=[2 / 3, 1 / 3])
    if image_index[0] == 0:
        image_flag = "FR"
        image_root = fr_root
        image_lst = fr_lst
        fr_flag = True
    else:
        image_flag = "FF"
        image_root = ff_root
        image_lst = ff_lst
        fr_flag = False

    pass


if __name__ == "__main__":
    cmd = "python test_video_swapsingle.py --crop_size 224 --use_mask --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path {} --video_path {} --output_path ./output/{} --temp_path ./temp_results"

    ff_root = "/disk1/generate_data/FF/10000"
    fr_root = "/disk1/generate_data/FR/front"
    fr_lst = os.listdir(fr_root)
    ff_lst = os.listdir(ff_root)

    # video_lst = ["assets/02.mp4", "assets/09.mp4", "assets/12.mp4"]
    vr_root = "/disk1/generate_data/VR/english"
    video_lst = os.listdir(vr_root)
    np.random.seed(0)

    num_count = 0
    while True:
        if num_count > 10000:
            break

        num_count += 1
        image_index = np.random.choice([0, 1], size=1, replace=False, p=[2 / 3, 1 / 3])
        if image_index[0] == 0:
            image_flag = "FR"
            image_root = fr_root
            image_lst = fr_lst
            fr_flag = True
        else:
            image_flag = "FF"
            image_root = ff_root
            image_lst = ff_lst
            fr_flag = False

        image_f = np.random.choice(image_lst, size=1, replace=False)[0]

        source_image_pth = os.path.join(image_root, image_f)
        driving_video_pth = os.path.join(
            vr_root, np.random.choice(video_lst, size=1, replace=False)[0]
        )

        out_name = (
            os.path.basename(image_f.split(".")[0])
            # os.path.basename(opt.source_image_pth).split(".")[0]
            + "_"
            + os.path.basename(driving_video_pth)
            # + os.path.basename(opt.driving_video_pth).split(".")[0]
        )
        if fr_flag:
            out_name = "FR_" + out_name
        else:
            out_name = "FF_" + out_name

        pic_path, video_path = source_image_pth, driving_video_pth

        if os.path.isfile(os.path.join("output", out_name)):
            print(out_name)
            continue
        try:
            os.system(cmd.format(pic_path, video_path, out_name))
        except:
            pass