#!/usr/bin/env bash
#
# MIT License
#
# Copyright (c) 2024 Mitsudori
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Use: ssim_optimizer.sh input.ext output.ext (to compare) bitrate_in_kbps (or none).
# The extension can be what ffmpeg was compiled with.
# The script will not work with downscale\upscale encode unless under comparative function (not implemented yet).
# Does not work with interlaced source yet.
# Depends: Bash, FFmpeg, FFprobe, Gnu AWK, Grep, bc, sed.
# Matroska may broke the results, because sometimes when FFmpeg encodes Matroska, inserts a frame plus.
# If no video correlation as param, the script will made a encode themselve.
# The execution under cygwin in Windows should not work for RAMdisk mounting, as you should do it manually.
# Not tested on Mac, BSD.
constantssim=0.979889
constantvmaf=96
qcmp=0.50
crf=16
ext=`echo "$1" | cut -d'.' -f2`
safeparams="-me_method hex -refs 8 -bf 8 -b_strategy 2 -trellis 2 -aq-mode 2 -x264-params partitions=all"
killparms="-me_method umh -me_range 32 -refs 16 -bf 16 -b_strategy 1 -trellis 2 -aq-mode 2 -x264-params partitions=all:subme=9:no-dct-decimate=1:no-fast-pskip=1"

if [ "$ext" = "yuv" ]; then 
	read -p "Enter size in widthxheight for YUV input." ryuv;
	yuvin="-s $ryuv -r 25 -pix_fmt yuv420p";
fi

if [ -z "$1" ]; then 
	printf "Arguments: ssim_optimizer.sh 'reference.mp4' 'correlation.mp4' 'initial_bitrate_int'\n\t\t\t\tAt least the reference needs to be insert." ; exit
fi 

if [ -z "$3" ] &&  [ ! -n "$2" ]; then
	read -p "Proceed to determine qcomp and qpfile (more precise encode)? (y/N)" q1;
	if [[ $q1 == [yY] || $q1 == [yY][eE][sS] ]]; then
		qcmp=`ffmpeg -i $1 -c:v libx264 -loglevel debug -qp 0 -g 1 -preset superfast -f null -y /dev/null 2>&1 | grep -oP "(?<=size=)\d+(?= b)" | gawk '!_[$0]++{z[++c]=$0}{l[NR]=$0}END{asort(z);for (a in l) for (b in z) if (l[a] == z[b]) {print b; break}}' | gawk '{q=$0;z[NR]=q;sum+=q}END{m=sum/NR;for (a in z) stdp+=(z[a]-m)**2;std=sqrt(stdp/NR); print std/m}'`;
		if [ -n "$yuvin" ]; then
			echo "QPfile function does not apply to yuv input."; else 
				ffprobe $1 -show_frames -print_format csv -show_entries stream=pict_type  | gawk -F"," '$19=="I"{print $20 " " $19 " -1" }' > qpfile_bruto.txt;
				telos=`gawk 'END{ print NR; }' qpfile_bruto.txt;`
		fi
		if [ $qcmp -gt 0.58 ]; then crf=$((crf-4)); fi
	fi
fi	

if [ -n "$1" ] && [ -n "$2" ] ; then 
	probe="-select_streams v -show_entries stream=width -of default=noprint_wrappers=1";
	a=`ffprobe -v error $probe "$2" | grep -Eo '[0-9]{1,4}' -`; b=`ffprobe -v error $probe "$1" | grep -Eo '[0-9]{1,4}' -`
	if [ "$a" != "$b" ]; then
		c=`ffprobe -v error $probe "$1" | grep -Eo '[0-9]{1,4}' -`
		parsedssim=`ffmpeg -hide_banner -i "$2" -i "$1" -filter_complex "[0:v]scale=$b:$c[v0r];[1:v][v0r] ssim" -f null - 2>&1 | grep -oP "(?<=All:).*(?= )"`
		if (( $(echo "$parsedssim < $constantssim" |bc -l) )); then	
			printf "\nYour correlation video is not optimized, with a round SSIM value of $parsedssim\n";
		else
			printf "\nYour correlation video is optimized, with a round SSIM value of $parsedssim\n";
		fi
	else
		parsedssim=`ffmpeg -hide_banner -i $1 -i $2 -lavfi ssim -f null - 2>&1 | grep -oP "(?<=All:).*(?= )"`;
		if (( $(echo "$parsedssim < $constantssim" |bc -l) )) ; then	
			printf "\nYour correlation video is not SSIM optimized, with a round SSIM value of $parsedssim";
		else
			printf "\nYour video is optimized, with a round SSIM value of $parsedssim";
		fi
		parsedvmaf=`ffmpeg -hide_banner -i $1 -i $2 -lavfi libvmaf -f null - 2>&1 | grep -oP "(?<=score: ).*"`;
		if  (( $(echo "$parsedvmaf < $constantvmaf" |bc -l) )) ; then	
			printf "\nYour correlation video is not VMAF optimized, with a round VMAF value of  $parsedvmaf";
		else
			printf "\nYour video is VMAF optimized, with a round VMAF value of  $parsedvmaf";
		fi
	exit
	fi
fi

encode_function() 
{
	read -p  "Procede to encode?" q2;
	if [[ $q2 == [yY] || $q2 == [yY][eE][sS] ]]; then
		fps=`ffprobe -hide_banner $1 2>&1 | grep -oP "(?<= )\d+(\.\d+)*(?= fps)" | gawk '{printf("%d\n",$1 + 0.5)}'`
		gop=$((10*$fps));
		scc=$((3*$fps));
		# for ((i=1;i<=telos;i++)); do     
			# flag=$((flag+1));
			# aframe`gawk -v var="$flag" 'NR==var{ print $1; }' qpfile_bruto.txt`
			# flag=$((flag+1));
			# zframe`gawk -v var="$flag" 'NR==var{ print $1; }' qpfile_bruto.txt`
		# done;
			# gawk 'NR==1{ print $1; }' qpfile_bruto.txt
		if [ $qcmp -gt 0.58 ]; then
			brate=`ffmpeg -i $1 -an -c:v libx264 -fastfirstpass 0 -pass 1  -qcomp $qcmp -g $gop -sc_threshold $scc -rc-lookahead $(($scc+$fps)) -crf $crf $killparms  -f null /dev/null | grep -oP "(?<=kb/s:).*"`
			ffmpeg -i $1 -an -c:v libx264 -pass 2  -qcomp $qcmp -g $gop -sc_threshold $scc -rc-lookahead $(($scc+$fps)) -b:v "$brate"k  -y encode.mp4
		else
			brate=`ffmpeg -i $1 -an -c:v libx264 -fastfirstpass 0 -pass 1  -qcomp $qcmp -g $gop -sc_threshold $scc -rc-lookahead $(($scc+$fps)) -crf $crf $safeparams  -f null /dev/null | grep -oP "(?<=kb/s:).*"`
			ffmpeg -i $1 -an -c:v libx264 -pass 2  -qcomp $qcmp -g $gop -sc_threshold $scc -rc-lookahead $(($scc+$fps)) -b:v "$brate"k  -y encode.mp4
		fi
	else
		exit;
	fi 
} 

encode_function