#!/usr/bin/perl -w
use Getopt::Std;
getopts("a:b:B:q:r:x:");
$search_after  = defined($opt_a) ? $opt_a : 2;
$search_before = defined($opt_b) ? $opt_b : 2;
$search_qp     = defined($opt_q) ? $opt_q : 2;
$reset_rate    = defined($opt_r) ? $opt_r : 2;
$bitrate       = defined($opt_B) ? $opt_B : 500;
$opt_x ||= "";
$opt_x =~ s/^x264\S*//;
$fps = 25;
$exe = "x264 --fps $fps -o /dev/null $opt_x";
$passfile = "x264_2pass.log";

sub write_qps(@) {
    open QPS, ">", "qps.txt" or die;
    for(0..$nframes-1) {
        print QPS "$_ $types[$_] $_[$_]\n";
    }
    close QPS;
}

sub run_rd($@) {
    my($cmd, @qps) = @_;
    my $bytes = 0;
    my $ssd = 0;
    if(@qps) {
        write_qps(@qps);
        if($memo_rd{"@qps"}) {
            return @{$memo_rd{"@qps"}};
        }
        $cmd .= " --qpfile qps.txt";
    }
    @encoder_log = `$cmd -v 2>&1`;
    $total_encodes++;
    for(@encoder_log) {
        /size=(\d+) bytes PSNR Y:(\d+\.\d+) U:(\d+\.\d+) V:(\d+\.\d+)/ or next;
        $bytes += $1;
        $ssd += 4*10**($2/-10)
                + 10**($3/-10)
                + 10**($4/-10);
    }
    $ssd /= 6;
    my $psnr;
    if(@qps) {
        $psnr = -10*log($ssd/$nframes)/log(10);
        #printf "$bytes, %.4f : @qps\n", $psnr;
        print ".";
        $memo_rd{"@qps"} = [$bytes, $ssd, $psnr];
    }
    return ($bytes, $ssd, $psnr);
}

sub run_dDdR($) {
    my($f) = @_;
    my @qps = @bqps;
    $qps[$f]--;
    if($qps[$f] < 0) {
        $dDdRs[$f] = 0;
        return;
    }
    my ($bytes, $ssd) = run_rd($exe, @qps);
    if($ssd > $bssd) {
        #print "nonmonotonic rd? frame $f qp $bqps[$f] -> $qps[$f]\n";
        $dDdRs[$f] = 0;
    } elsif($bytes > $target_bytes) {
        $dDdRs[$f] = 0;
    } else {
        if($bytes <= $bbytes) {
            $bytes = $bbytes+1;
        }
        $dDdRs[$f] = ($bssd-$ssd)/($bytes-$bbytes);
    }
}

sub max_idx(@) {
    my $m = $_[0];
    my $i = 0;
    for(1..$#_) {
        if($_[$_] > $m) {
            $m = $_[$_];
            $i = $_;
        }
    }
    return $i;
}

($bbytes, $bssd) = run_rd("$exe --crf=24 --pass 1");
$exe .= " --pass 2";

open PASS, "<", $passfile or die;
while(<PASS>){
    /in:(\d+).*type:(\w) q:(\d+)/ or next;
    $types[$1] = $2;
    $bqps[$1] = $3;
}
close PASS;
$nframes = scalar @types;
$target_bytes = $bitrate * 125 * $nframes / $fps;

my $qp_offset = $search_qp + int(log($bbytes / $target_bytes) * 6/log(2));
for(0..$nframes-1) {
    $bqps[$_] += $qp_offset;
}

($bbytes, $bssd) = run_rd($exe, @bqps);
for(0..$nframes-1) {
    run_dDdR($_);
}

$reset_counter = 0;
while(1) {
    my $f = max_idx(@dDdRs);
    last if($dDdRs[$f] <= 0);
    $bqps[$f]--;
    ($bbytes, $bssd, $psnr) = run_rd($exe, @bqps);
    if($bbytes >= $target_bytes) {
        $bqps[$f]++;
        $dDdRs[$f] = 0;
        next;
    }
    printf "\n%.3f, %.4f : @bqps ", $bbytes*$fps/$nframes/125, $psnr;
    foreach $j ($f - $search_before .. $f + $search_after) {
        next if($j < 0);
        last if($j >= $nframes);
        run_dDdR($j);
    }

    $reset_counter++;
    if($reset_counter > $nframes * $reset_rate) {
        for(0..$nframes-1) {
            run_dDdR($_);
        }
        $reset_counter = 0;
    }
}

%memo_rd = ();
($bbytes, $bssd, $psnr) = run_rd("$exe -o out.264", @bqps);
printf "\n%.3f, %.4f : @bqps\n", $bbytes*$fps/$nframes/125, $psnr;
printf "final: $bbytes/$target_bytes, psnr=%.4f, encodes=$total_encodes\n", $psnr;
open LOG, ">", "out.log" or die;
print LOG @encoder_log;
close LOG;