# Code adapted from "TSM: Temporal Shift Module for Efficient Video Understanding"
# Ji Lin*, Chuang Gan, Song Han

import os
import threading
import pdb

NUM_THREADS = 12
VIDEO_ROOT = './data/smt-smt-V2/smt-smt-V2-videos/20bn-something-something-v2/'
FRAME_ROOT = './data/smt-smt-V2/smt-smt-V2-frames/'


def split(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i + n]


def extract(video, tmpl='%06d.jpg'):
    cmd = 'ffmpeg -i \"{}/{}\" -threads 1 -vf scale=-1:256 -q:v 0 \"{}/{}/%06d.jpg\"'.format(VIDEO_ROOT, video,
                                                                                             FRAME_ROOT, video[:-5])
    os.system(cmd)


def target(video_list):
    for video in video_list:
        video_path = os.path.join(FRAME_ROOT, video[:-5])
        if not os.path.exists(video_path):
            os.makedirs(os.path.join(FRAME_ROOT, video[:-5]))
            extract(video)
        else:
            dir_files = os.listdir(os.path.join(FRAME_ROOT, video[:-5]))
            if len(dir_files) <= 10 and len(dir_files) != 0:
                print(f'folder {video} has only {len(dir_files)} frames')
                extract(video)



if __name__ == '__main__':
    if not os.path.exists(VIDEO_ROOT):
        raise ValueError('Please download videos and set VIDEO_ROOT variable.')
    if not os.path.exists(FRAME_ROOT):
        os.makedirs(FRAME_ROOT)

    video_list = os.listdir(VIDEO_ROOT)
    splits = list(split(video_list, NUM_THREADS))

    threads = []
    for i, split in enumerate(splits):
        thread = threading.Thread(target=target, args=(split,))
        thread.start()
        threads.append(thread)

    for thread in threads:
        thread.join()