"""
Index the official Kaggle training dataset and prepares a train and validation set based on folders

Video Face Manipulation Detection Through Ensemble of CNNs

Image and Sound Processing Lab - Politecnico di Milano

Nicolò Bonettini
Edoardo Daniele Cannas
Sara Mandelli
Luca Bondi
Paolo Bestagini
"""
import sys
import argparse
from multiprocessing import Pool
from pathlib import Path

import numpy as np
import pandas as pd
from tqdm import tqdm

from isplutils.utils import extract_meta_av


def parse_args(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', type=Path, help='Source dir', required=True)
    parser.add_argument('--videodataset', type=Path, default='data/dfdc_videos.pkl',
                        help='Path to save the videos DataFrame')
    parser.add_argument('--batch', type=int, help='Batch size', default=64)

    return parser.parse_args(argv)


def main(argv):
    ## Parameters parsing
    args = parse_args(argv)
    source_dir: Path = args.source
    videodataset_path: Path = args.videodataset
    batch_size: int = args.batch

    ## DataFrame
    if videodataset_path.exists():
        print('Loading video DataFrame')
        df_videos = pd.read_pickle(videodataset_path)
    else:
        print('Creating video DataFrame')

        # Create ouptut folder
        videodataset_path.parent.mkdir(parents=True, exist_ok=True)

        # Index
        df_train_list = list()
        for idx, json_path in enumerate(tqdm(sorted(source_dir.rglob('metadata.json')), desc='Indexing')):
            df_tmp = pd.read_json(json_path, orient='index')
            df_tmp['path'] = df_tmp.index.map(
                lambda x: str(json_path.parent.relative_to(source_dir).joinpath(x)))
            df_tmp['folder'] = int(str(json_path.parts[-2]).split('_')[-1])
            df_train_list.append(df_tmp)
        df_videos = pd.concat(df_train_list, axis=0, verify_integrity=True)

        # Save space
        del df_videos['split']
        df_videos['label'] = df_videos['label'] == 'FAKE'
        df_videos['original'] = df_videos['original'].astype('category')
        df_videos['folder'] = df_videos['folder'].astype(np.uint8)

        # Collect metadata
        paths_arr = np.asarray(df_videos.path.map(lambda x: str(source_dir.joinpath(x))))
        height_list = []
        width_list = []
        frames_list = []
        with Pool() as pool:
            for batch_idx0 in tqdm(np.arange(start=0, stop=len(df_videos), step=batch_size), desc='Metadata'):
                batch_res = pool.map(extract_meta_av, paths_arr[batch_idx0:batch_idx0 + batch_size])
                for res in batch_res:
                    height_list.append(res[0])
                    width_list.append(res[1])
                    frames_list.append(res[2])

        df_videos['height'] = np.asarray(height_list, dtype=np.uint16)
        df_videos['width'] = np.asarray(width_list, dtype=np.uint16)
        df_videos['frames'] = np.asarray(frames_list, dtype=np.uint16)

        print('Saving video DataFrame to {}'.format(videodataset_path))
        df_videos.to_pickle(str(videodataset_path))

    print('Real videos: {:d}'.format(sum(df_videos['label'] == 0)))
    print('Fake videos: {:d}'.format(sum(df_videos['label'] == 1)))


if __name__ == '__main__':
    main(sys.argv[1:])