#!/usr/bin/env python3

import os

import time

import csv

import ffmpy

from multiprocessing import cpu_count

from collections import OrderedDict

import json

# always round timestamp to integer
def now():
    return int(time.time())


def write_line(
    codec: str,
    crf: str,
    preset: str,
    infile: str,
    outfilesize: float,
    enctime: int,
    vmafmean: float,
    vmafmin: float,
):
    """
    Write line to data csv

    Parameters:
        codec (str): Codec used
        crf (str): CRF used
        preset (str): Preset used
        infile (str): Input file name
        outfilesize (float): Size of output file
        enctime (int): Time to encode
        vmafmean (float): Mean VMAF score
        vmafmin (float): Min VMAF score
    """
    with open(datafile, "a", newline="") as file:
        write = csv.writer(file)
        write.writerow(
            (
                codec,
                crf,
                preset,
                infile,
                outfilesize,
                enctime,
                vmafmean,
                vmafmin,
            )
        )


if __name__ == "__main__":
    encoding = {
        "libx264": {
            "crf": [15, 20, 25, 30, 35],
            "presets": ["fast", "medium", "slow", "veryslow"],
        },
        "libx265": {
            "crf": [15, 20, 25, 30, 35],
            "presets": ["fast", "medium", "slow", "veryslow"],
        },
        "libaom-av1": {"crf": [20, 25, 30, 35, 40], "presets": [0, 2, 4, 6]},
        "libsvtav1": {"crf": [20, 25, 30, 35, 40], "presets": [0, 4, 8, 12]},
    }

    if not os.path.isdir("encoded"):
        os.mkdir("encoded")

    datafile = "data-" + str(now()) + ".csv"

    with open(datafile, "w", newline="") as file:
        write = csv.writer(file)
        write.writerow(
            (
                "Codec",
                "CRF",
                "Preset",
                "Input file",
                "Output file size (MiB)",
                "Encode time (s)",
                "VMAF Score (mean)",
                "VMAF Score (min)",
            )
        )

    for codec in encoding:
        for crf in encoding[codec]["crf"]:
            for preset in encoding[codec]["presets"]:
                outputfile = os.path.join(
                    "encoded",
                    (
                        "Sparks_in_Blender-codec_"
                        + codec
                        + "-crf_"
                        + str(crf)
                        + "-preset_"
                        + str(preset)
                        + ".mkv"
                    ),
                )

                # libaom needs additional options
                if codec == "libaom-av1":
                    ff = ffmpy.FFmpeg(
                        inputs={"Sparks_in_Blender.webm": None},
                        outputs={
                            outputfile: "-c:v {videocodec} -crf {crf} -b:v 0 -cpu-used {preset} -row-mt 1 -tiles 2x2 -g 240 -map 0:v:0 ".format(
                                videocodec=codec,
                                crf=crf,
                                preset=preset,
                            )
                        },
                    )
                else:
                    ff = ffmpy.FFmpeg(
                        inputs={"Sparks_in_Blender.webm": None},
                        outputs={
                            outputfile: "-c:v {videocodec} -crf {crf} -preset {preset} -g 240 -map 0:v:0 ".format(
                                videocodec=codec,
                                crf=crf,
                                preset=preset,
                            )
                        },
                    )

                starttime = now()
                ff.run()
                endtime = now()
                difftime = int(endtime - starttime)

                outputfilesize = os.path.getsize(outputfile) / 1024 / 1024

                ffvmaf = ffmpy.FFmpeg(
                    inputs=OrderedDict(
                        [(outputfile, None), ("Sparks_in_Blender.webm", None)]
                    ),
                    outputs={
                        "-": "-filter_complex libvmaf=log_fmt=json:n_threads={cputhreads}:log_path=vmaf.json -f null".format(
                            cputhreads=cpu_count()
                        )
                    },
                )

                ffvmaf.run()

                with open("vmaf.json", "r") as file:
                    vmaf = json.load(file)

                write_line(
                    codec=codec,
                    crf=crf,
                    preset=preset,
                    infile="Sparks_in_Blender.webm",
                    outfilesize=outputfilesize,
                    enctime=difftime,
                    vmafmean=vmaf["pooled_metrics"]["vmaf"]["mean"],
                    vmafmin=vmaf["pooled_metrics"]["vmaf"]["min"],
                )

                os.remove(outputfile)