#!/bin/env python

# neural style transfer video generator
# 2020-02-11 Oliver Meckmann

# meant to be run from the root directory of this project:
# https://github.com/ProGamerGov/neural-style-pt

# opticalflow.py from this project should also be copied over to the
# neural-style-pt root directory
# https://github.com/larspars/neural-style-video


from gooey import Gooey, GooeyParser
import subprocess
import os
import glob

@Gooey(default_size=(800, 600), terminal_font_color="#FFFFFF", \
        header_bg_color="#000000")
def init():
    parser = GooeyParser(description = \
            "nstvid - Neural Style Transfer Video Creator")

    parser.add_argument("style_image", \
            help="Style image(s) separated by comma", \
            widget="FileChooser")
    parser.add_argument("input_video", help="Source video", \
            widget="FileChooser", default="/home/pux/tmp/nstest/in.mp4")
    parser.add_argument("output_dir", help="Where your video is", \
            widget="DirChooser", default="/home/pux/tmp/nstest")
    parser.add_argument("fps", help="Framerate", default=12)
    parser.add_argument("image_size", help="Pixels", default=384)
    parser.add_argument("style_weight", default=1000)
    parser.add_argument("content_weight", default=5)

    params = parser.parse_args()
    return params

def main():
    print("Running ffmpeg extract images")
    command = [
            "ffmpeg", "-i", r"" + params.input_video,
            "-r", r"" + params.fps, "-f", "image2",
            r"" + params.output_dir + os.sep + "image-%5d.jpg"]
    subprocess.Popen(command).wait()

    print("Running ffmpeg extract audio")
    command = [
            "ffmpeg", "-y", "-i", r"" + params.input_video,
            r"" + params.output_dir + os.sep + "rawaudio.wav"]
    subprocess.Popen(command).wait()

    left = len(params.output_dir)
    flist = sorted(glob.glob(params.output_dir + os.sep + "image-*.jpg"))
    for fn in flist:
        num = fn[left+len(os.linesep)+6:-4]
        print("Processing loop: new iteration")
        print("Num: " + num + ", input file: " + fn)

        command = [
                "python", "neural_style.py",
                "-style_image", r"" + params.style_image,
                "-content_image", r"" + params.output_dir + os.sep + "image-" \
                + str(num) + ".jpg",
                "-output_image", r"" + params.output_dir + os.sep \
                + "generated-" + str(num) + ".jpg",
                "-image_size", params.image_size,
                "-style_weight", params.style_weight,
                "-content_weight", params.content_weight,
                "-cudnn_autotune",
                "-tv_weight", "000085", "-num_iterations", "1000",
                "-print_iter", "50",
                "-init", "image", "-optimizer", "lbfgs",
                "-learning_rate", "1.0",
                "-lbfgs_num_correction", "100",
                "-style_scale", "1.0", 
                "-pooling", "max",
                "-backend", "cudnn",
                "-seed", "100",
                "-content_layers", "relu1_1,relu2_1,relu3_1,relu4_1,relu4_2," \
                + "relu5_1",
                "-style_layers", "relu3_1,relu4_1,relu4_2,relu5_1",
                "-multidevice_strategy", "4,7,29"]
        subprocess.Popen(command).wait()

        if fn is not flist[-1]:
            print("Running opticalflow: num: " + num)
            command = [
                    "python", "opticalflow.py",
                    "--input_1", r"" + params.output_dir \
                    + os.sep + "image-" + str(num) + ".jpg",
                    "--input_1_styled", r"" + params.output_dir \
                    + os.sep + "generated-" + str(num) + ".jpg",
                    "--input_2", r"" + params.output_dir + os.sep + "image-" \
                    + str(int(num)+1).zfill(5) + ".jpg",
                    "--output", r"" + params.output_dir \
                    + os.sep + "blended-" + str(num) + ".jpg",
                    "--alpha", "0.55",
                    "--flowblend", "False"]
            subprocess.Popen(command).wait()
        else:
            print("Last element: skipping opticalflow")

    print("Running ffmpeg merge")
    command = [
            "ffmpeg", "-y", "-framerate", params.fps,
            "-i", r"" + params.output_dir + os.sep + "blended-%05d.jpg",
            "-i", r"" + params.output_dir + os.sep + "rawaudio.wav",
            "-c:v", "libx264", "-preset", "veryslow", "-qp", "0",
            "-pix_fmt", "yuv420p", r""
            + params.output_dir + os.sep + "out.mp4"]
    subprocess.Popen(command).wait()

if __name__ == '__main__':
    params = init()
    print(params.style_image)
    main()