import os
import argparse
from itertools import zip_longest

def img2mp4(div, trial):
    log_path = os.path.join(args.dir, div, trial)
    print(log_path)
    output_mp4 = os.path.join('data/videos', args.dir, f'{div}_{trial}.mp4')
    # if os.path.isfile(output_mp4):
    #     print(f'{output_mp4} already exists!')
    #     return
    cmd = 'ffmpeg -framerate %d -i ' % args.fps + log_path + '/plt/vis_step_%06d.png ' + output_mp4
    print(cmd)
    os.system(f"ffmpeg -pattern_type glob -framerate {args.fps} -i '{log_path}/plt/vis_step_*.png' {output_mp4}")

def img2gif(div, trial):
    log_path = os.path.join(args.dir, div, trial)
    print(log_path)
    output_gif = os.path.join('data/videos', args.dir, f'{div}_{trial}.gif')
    import imageio
    images = []
    filenames = os.listdir(os.path.join(log_path, 'plt'))
    # print(filenames)
    for filename in filenames:
        print(os.path.join(log_path, 'plt', filename))
        images.append(imageio.imread(os.path.join(log_path, 'plt', filename)))
    imageio.mimsave(output_gif, images, duration=0.2)

    if os.path.isfile(output_mp4):
        print(f'{output_mp4} already exists!')
        return
    os.system(f"ffmpeg -pattern_type glob -framerate {args.fps} -i '{log_path}/plt/vis_step_*.png' {output_mp4}")
   
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dir", type=str)
    parser.add_argument('--fps', type=int, default=5)
    args = parser.parse_args()

    divs = sorted(os.listdir(args.dir))
    log_lists = [sorted(os.listdir(os.path.join(args.dir, div))) for div in divs]
    os.makedirs(os.path.join('data/videos', args.dir), exist_ok=True)

    for idx, trials in enumerate(zip_longest(*log_lists)):
        for div, trial in zip(divs, trials):
            if trial is None:
                continue
            file_path = os.path.join(args.dir, div, trial)
            # print(file_path)
            img2mp4(div, trial)
            # img2gif(div, trial)
