from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import sys
import argparse
import json
import subprocess
import logging
import numpy as np
import matplotlib.image as mpimg
from functools import partial
from multiprocessing.pool import ThreadPool

from video_process_lib.ffmpeg import copy_ffmpeg_bin


FORMAT = '[%(levelname)s: %(filename)s: %(lineno)4d]: %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT, stream=sys.stdout)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


def process_video(v_data, input_video_dir, output_video_dir):
    cwd = '/tmp'
    fps = 25
    border_thickness = 0.03
    youtubeID = v_data[0]
    video_fn = os.path.join(input_video_dir, youtubeID + '.mp4')
    FFPROBE_BIN = '{}/ffmpeg/ffmpeg-2.6.1/bin/ffprobe'.format(cwd)
    cmd = '{} -v quiet -print_format json -show_format -show_streams {}'.format(
        FFPROBE_BIN, video_fn
    )
    logger.info('ffprobe cmd: {}'.format(cmd))
    split_cmd = cmd.split(' ')

    height, width = 0, 0
    try:
        video_info = json.loads(subprocess.check_output(split_cmd))
        for stream in video_info['streams']:
            if 'height' in stream and 'width' in stream:
                height = stream['height']
                width = stream['width']
                logger.info(
                    '{} | Get video height/width: {}/{}'.format(
                        video_fn, height, width)
                )
                break
    except BaseException as e:
        logger.warning(
            '{} | Fail to get height and width. Error: {}'.format(video_fn, e))
        return
    # decode video into images
    FFMPEG_BIN = '{}/ffmpeg/ffmpeg-2.6.1/bin/ffmpeg'.format(cwd)
    cmd = '{} -i {} -r {} /tmp/{}_%05d.png'.format(
        FFMPEG_BIN, video_fn, fps, youtubeID
    )
    logger.info('frame extraction cmd: {}'.format(cmd))
    ret_val = subprocess.call(cmd, shell=True)
    if ret_val == 0:
        print('video succeed: {}'.format(video_fn))
    else:
        print('video fails: {}'.format(video_fn))

    if len(v_data[1]["annotation"]) == 0:
        logger.info("Completely negative video: {}. Return !".format(video_fn))
        return

    # modify those images with ongoing actions so that border top border color is
    green = np.asarray([0., 1., 0.])
    for seg in v_data[1]["annotation"]:
        start, end = seg[0], seg[1]
        start_img = int(round(start * fps)) + 1
        end_img = int(round(end * fps)) + 1
        for img_id in range(start_img, end_img + 1):
            img_fn = "/tmp/{}_{:05d}.png".format(youtubeID, img_id)
            if os.path.exists(img_fn):
                im = mpimg.imread(img_fn)
                logger.info("img: {} shape: {}".format(img_fn, im.shape))
                height, width = im.shape[0], im.shape[1]

                border_height_pix = int(round(border_thickness * height))
                border_width_pix = int(round(border_thickness * width))

                green_border = np.tile(green, [border_height_pix, width, 1])
                im[:border_height_pix, :, :] = green_border
                im[-border_height_pix:, :, :] = green_border

                green_border = np.tile(green, [height, border_width_pix, 1])
                im[:, :border_width_pix, :] = green_border
                im[:, -border_width_pix:, :] = green_border

                logger.info("Modify {}".format(img_fn))
                mpimg.imsave(img_fn, im)

    # assemble images into video and save it
    v_height = (height // 2) * 2
    v_width = (width // 2) * 2
    out_fn = os.path.join(output_video_dir, '{}.mp4'.format(youtubeID))
    cmd = '{} -framerate {} -i /tmp/{}_%05d.png -vcodec libx264  -pix_fmt yuv420p -s {}x{} {}'.format(
        FFMPEG_BIN, fps, youtubeID, v_width, v_height, out_fn
    )
    logger.info('video assembling cmd: {}'.format(cmd))
    ret_val = subprocess.call(cmd, shell=True)
    if ret_val == 0:
        print('video assembling succeed: {}'.format(out_fn))
    else:
        print('video assembling fails: {}'.format(out_fn))

    # remove frame images
    cmd = "rm -f /tmp/{}*.png".format(youtubeID)
    ret_val = subprocess.call(cmd, shell=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_video_dir',
                        default='/mnt/vol/gfsai-oregon/ai-group/users/hangzhao'
                        '/hac650k_export0828_val_videos')
    parser.add_argument('--input_annotation',
                        default='/mnt/vol/gfsai-oregon/ai-group/users/hangzhao'
                        '/annotation/full_annotation/full_annotation_v1_val.json')
    parser.add_argument('--output_video_dir',
                        default='/mnt/vol/gfsai-oregon/ai-group/users/hangzhao'
                        '/annotation/full_annotation/vis')
    parser.add_argument('--num_worker', default=36, type=int)
    parser.add_argument('--shard_id', default=0, type=int)
    parser.add_argument('--num_shard', default=2000, type=int)

    args = parser.parse_args()

    logger.info("Extract video meta info")
    copy_ffmpeg_bin(
        '/mnt/vol/gfsai-oregon/ai-group/users/hangzhao/ffmpeg.tar.gz')

    with open(args.input_annotation) as fp:
        anno = json.load(fp)
    logger.info("{} videos".format(len(anno.keys())))

    if not os.path.exists(args.output_video_dir):
        os.mkdir(args.output_video_dir)
        os.chmod(args.output_video_dir, 0o777)

    list_anno = list(anno.items())[:5]

    list_anno = []
    for item in list(anno.items()):
        if hash(item[0]) % args.num_shard == args.shard_id:
            list_anno.append(item)
    logger.info("{} videos in shard {}/{}".format(
        len(list_anno), args.shard_id, args.num_shard))

    pool = ThreadPool(args.num_worker)
    handles = pool.map(
        partial(
            process_video,
            input_video_dir=args.input_video_dir,
            output_video_dir=args.output_video_dir,
        ),
        list_anno)