#!/bin/bash -e
# Encode a TV programme recorded by MythTV into a more convenient format.
# This uses the heuristics suggested in the MPlayer manual to figure out the
# output frame size.
# Adam Sampson <ats@offog.org>

ifile="not-set"
ofile="not-set"
tsize=$((350 * 1024 * 1024))
daspect=$(python -c "print 16.0/9")
cq=0.21
aformat="none"
aquality=3
cropw=""
croph=""
tempdir="/var/tmp"
decopts=""
encopts=""

mplayer="mplayer"
mencoder="mencoder"

usage () {
	echo "Usage: encode-prog [OPTIONS] INPUT-FILE"
	echo
	echo "-o FILE    Write output to FILE (default INPUT-FILE.mkv)"
	echo "-s SIZE    Aim for output files of size SIZE megabytes (default 350)"
	echo "-q QUALITY Use quality QUALITY for video (default 0.21)"
	echo "-O         Reencode audio to Ogg Vorbis (default leave as MP2)"
	echo "-Q QUALITY Use quality QUALITY for reencoded audio (default 3)"
	echo "-a ASPECT  Input files have aspect ratio ASPECT (default 16/9)"
	echo "-W WIDTH   Crop to WIDTH pixels wide (e.g. 540 for BBC 4:3)"
	echo "-H HEIGHT  Crop to HEIGHT pixels tall (e.g. 432 for Film4 2.39:1)"
	echo "-t TEMPDIR Use TEMPDIR as temporary directory (default /var/tmp)"
	echo "-d EXTRA   Give extra options EXTRA to mplayer/mencoder (default none)"
	echo "-e EXTRA   Give extra options EXTRA to mencoder (default none)"
}

while getopts "o:s:q:OQ:a:W:H:t:d:e:" c ; do
	case $c in
	o)
		ofile="$OPTARG"
		;;
	s)
		tsize=$(($OPTARG * 1024 * 1024))
		;;
	q)
		cq="$OPTARG"
		;;
	O)
		aformat="vorbis"
		;;
	Q)
		aquality="$OPTARG"
		;;
	a)
		daspect=$(python -c "print $(echo "$OPTARG" | sed 's,/,.0/,')")
		;;
	W)
		cropw="$OPTARG"
		;;
	H)
		croph="$OPTARG"
		;;
	t)
		tempdir="$OPTARG"
		;;
	d)
		decopts="$OPTARG"
		;;
	e)
		encopts="$OPTARG"
		;;
	\?)
		usage
		exit 0
		;;
	esac
done
shift `expr $OPTIND - 1`

if [ "$1" = "" ] ; then
	usage
	exit 1
fi
ifile="$1"
if [ "$ofile" = "not-set" ] ; then
	ofile="$ifile.mkv"
fi

if [ "$aformat" = "vorbis" ] ; then
	afile="$tempdir/audio.ogg"
else
	afile="$tempdir/audio.mp2"
fi

clean () {
	rm -f $tempdir/audio.wav audio.log $afile lavc_stats.txt divx2pass.log $tempdir/video.avi
}

clean

isize=$(echo q | $mplayer $decopts -ao null -vo null -v "$ifile" 2>&1 | grep ^VIDEO: | awk '{print $3}')
iwidth=$(echo $isize | sed 's,x.*,,')
if [ "$cropw" != "" ] ; then
	iwidth="$cropw"
fi
iheight=$(echo $isize | sed 's,.*x,,')
if [ "$croph" != "" ] ; then
	iheight="$croph"
fi

if [ "$aformat" = "vorbis" ] ; then
	$mplayer $decopts -vc null -vo null -ao pcm:fast:file=$tempdir/audio.wav "$ifile" | tee audio.log
	length=$(tr '\r' '\n' <audio.log | grep '^A:' | tail -1 | sed 's,.* V: *\([^ ]*\) .*,\1,')
	oggenc -q$aquality $tempdir/audio.wav
else
	$mencoder $decopts -o $afile -oac copy -ovc copy -of rawaudio "$ifile" | tee audio.log
	length=$(grep '^Pos:' audio.log | sed 's,.*Pos: *\([^ ]*\)s.*,\1,')
fi

asize=$(stat -c%s $afile)
vbitrate=$(python -c "print int((($tsize - $asize) * 8 / $length) / 1000)")
iaspect=$(python -c "print (1.0 * $iwidth) / $iheight")
oaspect=$(python -c "print $iwidth * (($daspect) / ($iaspect)) / $iheight")
oheight=$(python -c "import math; print int(math.sqrt(1000 * $vbitrate / 25 / ($oaspect) / $cq) / 16) * 16")
owidth=$(python -c "print int($oheight * ($oaspect) / 16) * 16")

filters="pp=hb:a/vb:a/dr:a/fd,scale=$owidth:$oheight,hqdn3d=2:1:2"
if [ "$cropw" != "" -o "$croph" != "" ] ; then
	filters="crop=${cropw:-$iwidth}:${croph:-$iheight},$filters"
fi
common="$decopts $encopts -vf $filters -ovc lavc -oac copy"
# These cause an explosion in ffmpeg for me at the moment...
# lavcopts="vcodec=mpeg4:vbitrate=$vbitrate:v4mv:mbd=2:trell:cmp=3:subcmp=3:mbcmp=3:autoaspect"
lavcopts="vcodec=mpeg4:vbitrate=$vbitrate:mbd=2:trell:mbcmp=3:autoaspect"

echo
echo "Source material is $iwidth x $iheight, length $length seconds"
echo "Encoding at video bitrate $vbitrate, size $owidth x $oheight"
echo "Encoding options: $common -lavcopts $lavcopts"
echo

$mencoder -o $tempdir/video.avi $common -lavcopts $lavcopts:vpass=1 "$ifile"
$mencoder -o $tempdir/video.avi $common -lavcopts $lavcopts:vpass=2 "$ifile"

rm -f "$ofile"
mkvmerge -o "$ofile" -A $tempdir/video.avi $afile

echo
echo "Output file:"
ls -lh "$ofile"
echo

clean