#!/usr/bin/env bash
#
# v2022-06-13
#
# I run it like this (using zsh syntax):
#
#   $ for video (~/Pictures/2021/**/*.mp4(ND.)); do ~/Downloads/av1-tests/convert-libaom.sh "$video"; done 
#
# Reference:
#   - https://ffmpeg.org/ffmpeg-codecs.html#libaom_002dav1
#   - https://trac.ffmpeg.org/wiki/Encode/AV1
#
# Changes:
# 2024-01-21: minor improvement to handling of output file name
# 2022-06-13: use _max_crf in mkvpropedit if we don't find an acceptable vmaf
# 2022-04-05: update VMAF log parsing
# 2021-10-21: I tested 2-pass and found it doesn't do anything for constant
#             quality mode. cpu-used 2 still takes ~5x longer than cpu-used 4
#             at all tiling levels. 4x4 tiles much faster than 2x1 and 2x2 on
#             all cpu-used levels, with quality less than 0.5 VMAF difference.
#             Using 4x4 creates a grid of 16 tiles which just seems wrong...
#             especially if you don't have many threads (I was using 12), so
#             I will default to 2x2 as a compromise.
# 2021-10-20: use tiles instead of columns and rows
# 2021-10-16: use tile-columns 2 and tile-rows 2 (oops, this is 4x4 tiles!)
# 2021-08-31: strip spaces and commas from VMAF log file name as well
# 2021-05-14: benchmarked two pass and file sizes are larger by .5-1MB
# 2021-04-29: make sure input file exists
# 2021-03-07: benchmarked 10-bit and it's only .5 or less improvment to VMAF
# 2021-02-22: allow overriding threads
# 2021-02-15: fix handling of apostrophes in ffmpeg's log_path
# 2021-01-27: round VMAF before comparing, simplifies code a bit
# 2021-01-25: stop trying lower CRFs if VMAF is not likely to improve
# 2021-01-19: embed encoding parameters in webm metadata with mkvpropedit
# 2021-01-18: use 1-pass encoding (2-pass is only for trying to hit a target
#             bitrate!)
# 2021-01-18: allow overriding variables
# 2021-01-17: detect if video already processed
# 2021-01-08: use cpu-used 4
# 2021-01-07: use tile-columns 2 and tile-rows 2 with cpu-used 4
# 2021-01-07: use tile-columns 2 and tile-rows 2 with cpu-used 3
# 2021-01-07: use -g 300 (30fps x 10)

# exit on first error
set -o errexit

if [[ -z "$1" ]]; then
    echo "No input file specified."

    exit 1
elif [[ ! -r "$1" ]]; then
    echo "Input file missing or unreadable: $1"

    exit 1
fi

INPUT_FILE_BASENAME=$(basename -- "$1")
INPUT_FILE_EXTENSION=${INPUT_FILE_BASENAME##*.}
INPUT_FILE_DIRNAME=$(dirname -- "$1")

# Check if an output file name was specified (like if we are calling from the
# benchmarking script, in which case we want to exit as soon as possible). If
# not then we can continue with a simple filename based on the input file's.
if [[ -z $_output_file ]]; then
    # Strip file extension from end of string and replace with webm
    _output_file="${INPUT_FILE_BASENAME%$INPUT_FILE_EXTENSION}webm";
    _benchmark_mode="false"
else
    _benchmark_mode="true"
fi

# aq-mode not recommended in AV1 yet
_aq_mode=${_aq_mode:-0}
# 2021-01-08 (libaom 2.0.1): realistic range between 3 and 5, lower takes *much*
# longer with very little boost to VMAF. In my experience, cpu-used 2 is *three*
# times slower than cpu-used 4 at the same CRF for only a 0.5% boost in VMAF. 6
# is the same as 5 in all my tests *shrug*.
_cpu_used=${_cpu_used:-4}
# 2021-10-20 -tile-columns and -tile-rows are for compatibility with libvpx. The
# new option for libaom is -tiles. Convert any columns/rows to tiles using 2^n.
if [[ ! -z $_tile_columns || ! -z $_tile_rows ]]; then
    _tile_columns=${_tile_columns:-1}
    _tile_rows=${_tile_rows:-1}

    # Compute 2^n tile columns and tile rows (the power operator is ** in bash)
    _tiles="$((2 ** $_tile_columns))x$((2 ** $_tile_rows))"
else
    _tiles=${_tiles:-2x2}
fi

# 2021-02-22 use nproc number of threads unless threads is already set.
_nproc=$(nproc)
_threads=${_threads:-$_nproc}

# Range of CRFs to try, where higher is faster, but lower quality. Based on my
# testing 52 to 40 should cover most crappy phone videos.
_crf_max=${_crf_max:-52}
_crf_min=${_crf_min:-40}

_libaom_version=$(pacman -Qi aom | grep Version | awk '{print $3}')

# For grainy mobile phone videos 90 is fine
_target_vmaf=${_target_vmaf:-90}
# Don't bother trying lower CRF values if the CRF score is more than five points
# above the target. In my experience each successively lower CRF step gives you
# ~1 VMAF point. In these cases it isn't likely we'll ever reach the target, so
# we might as well just settle on the current CRF straight away.
_target_vmaf_threshold=${_target_vmaf_threshold:-5}

_acceptable_output_found="no"

# Change to the input file's directory
pushd "$INPUT_FILE_DIRNAME" >/dev/null

# We want the highest VMAF score possible with the highest CRF possible. Start
# with high CRF first to see if we can get an acceptable VMAF score as soon as
# possible. Step through values by 2.
for _crf in $(seq $_crf_max -2 $_crf_min); do
    _processed=$(find . -maxdepth 1 -type f -iname "$_output_file" | wc -l)
    if [[ $_processed -gt 0 ]]; then
        echo "${INPUT_FILE_BASENAME}: already processed"

        exit 0
    fi

    if [[ $_benchmark_mode == "false" ]]; then
        echo "Processing ${INPUT_FILE_BASENAME} with libaom $_libaom_version CRF ${_crf}..."
    fi

    chrt -b 0 nice ffmpeg -hide_banner -y -i "$INPUT_FILE_BASENAME" -c:v libaom-av1 -b:v 0 -crf $_crf \
         -aq-mode $_aq_mode -c:a libopus -b:a 16k \
         -sc_threshold 0 -cpu-used $_cpu_used \
         -tiles $_tiles -row-mt 1 \
         -auto-alt-ref 1 -lag-in-frames 25 \
         -g 300 -threads $_threads \
         -f webm "$_output_file" 2>/dev/null

    # Return quickly if we are in benchmark mode so the benchmark script can get
    # an accurate time and compute its own VMAF score.
    if [[ $_benchmark_mode == "true" ]]; then
        exit 0
    fi

    _vmaf_log="${_output_file/.*/}.log"

    # strip apostrophes, spaces, and commas (if any) from log file name because
    # ffmpeg's log_path doesn't seem to be able to handle them. The // performs
    # a global replace.
    _vmaf_log="${_vmaf_log//\'}"
    _vmaf_log="${_vmaf_log//\,}"
    _vmaf_log="${_vmaf_log// }"

    # Get VMAF score with harmonic mean to emphasize small outliers
    # See: https://netflixtechblog.com/vmaf-the-journey-continues-44b51ee9ed12
    chrt -b 0 nice ffmpeg -hide_banner -y -i "$_output_file" -i "$INPUT_FILE_BASENAME" \
         -filter_complex "libvmaf=pool=harmonic_mean:log_path=${_vmaf_log}:log_fmt=json" \
         -f null - 2>/dev/null

    _vmaf_score=$(jq '.["pooled_metrics"]["vmaf"]["harmonic_mean"]' "$_vmaf_log")
    # bash can't do floating point so we use bc
    # See: https://stackoverflow.com/questions/8654051/how-to-compare-two-floating-point-numbers-in-bash
    _vmaf_score_round=$(printf %.$2f $(echo "scale=0;(((10^0)*$_vmaf_score)+0.5)/(10^0)" | bc))

    rm "$_vmaf_log"

    # Check if rounded VMAF score is >= target VMAF
    if [[ $_vmaf_score_round -ge $_target_vmaf ]]; then
        printf "$INPUT_FILE_BASENAME: acceptable VMAF (%.3f) at AV1 CRF ${_crf}.\n" $_vmaf_score

        # Set the title with information about the encoding in the MKV title
        mkvpropedit --edit info --set "title=${INPUT_FILE_BASENAME/.*/}-libaom_${_libaom_version}-crf${_crf}-cpu${_cpu_used}-${_tiles}-tiles-vmaf_${_vmaf_score}" "$_output_file" >/dev/null

        _acceptable_output_found="yes"

        # Break from the for loop because we have an acceptable output
        break
    # Check if the VMAF score is anywhere near our target, otherwise just keep
    # current output.
    elif (($_target_vmaf - $_vmaf_score_round >= $_target_vmaf_threshold)); then
        printf "$INPUT_FILE_BASENAME: unacceptable VMAF (%.3f) at AV1 CRF ${_crf}, unlikely to reach target soon.\n" $_vmaf_score

        _acceptable_output_found="yes"

        # Break from the for loop and keep the current output, even though it's
        # unacceptable.
        break
    else
        printf "$INPUT_FILE_BASENAME: unacceptable VMAF (%.3f) at AV1 CRF ${_crf}, continuing.\n" $_vmaf_score

        # Clean up the unacceptable output
        rm "$_output_file"
    fi
done

# If we finished going over all the CRFs and still didn't get an acceptable VMAF
# score then the video is probably just really low quality so let's just convert
# using our max CRF and be done with it.
if [[ $_acceptable_output_found == "no" ]]; then
    echo "${INPUT_FILE_BASENAME}: no acceptable output found, settling on AV1 CRF ${_crf_max}."

    chrt -b 0 nice ffmpeg -hide_banner -y -i "$INPUT_FILE_BASENAME" -c:v libaom-av1 -b:v 0 -crf $_crf_max \
         -aq-mode $_aq_mode -c:a libopus -b:a 16k \
         -sc_threshold 0 -cpu-used $_cpu_used \
         -tiles $_tiles -row-mt 1 \
         -auto-alt-ref 1 -lag-in-frames 25 \
         -g 300 -threads $_threads \
         -f webm "$_output_file" 2>/dev/null

        # Set the title with information about the encoding in the MKV title
        mkvpropedit --edit info --set "title=${INPUT_FILE_BASENAME/.*/}-libaom_${_libaom_version}-crf${_max_crf}-cpu${_cpu_used}-${_tiles}-tiles" "$_output_file" >/dev/null
fi

# Change back to our starting directory
popd >/dev/null

exit 0