/*
 * $Id: psnrcalc.c,v 1.5 2005/03/02 07:36:52 hzoli Exp hzoli $ 
 * Licence: GNU GPL v2
 * Copyright: Zoltan Hidvegi <mplayer -a- hzoli -d- com>
 * Read and plot lavc or x264 psnr files
 * Compile: gcc -Wall -lm -O2 -o psnrcalc psnrcalc.c
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <signal.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <math.h>
#include <limits.h>

#if !defined(INFINITY) && defined(HUGE_VAL)
#define INFINITY HUGE_VAL
#endif

const char *usage =
"psnrcalc [OPTIONS] [FILE]\n"
"  [FILE]	The PSNR file, read from stdin if not specified\n"
"  -fps <float>	Frames/sec for edl and bitrate calculations [23.976]\n"
"  -edl <file>	Skip edl blocks\n"
"  -t<IPB>	Use only the given frame types.\n"
"  -p[num]	GNUPlot PSNR dist. on <num> sample intervals [100]\n"
"  -x<float>	Exponent for average calculatin. Higher values lower\n"
"		average of more dispersed distributions [1].\n"
"  -i<float>	Ignore frames with PSNR greater than [60.0]\n"
"  -l<float>	Ignore frames with PSNR less than [0.0]\n"
"  -<num>	Use only the first <num> frames.\n"
"  -h		Show this help message\n";

typedef struct Sample {
    float r;
    unsigned size;
} Sample;

static double avx;

static inline double
invpsnr(double x)
{
    if (x == INFINITY)
	return 0;
    else
	return exp(x * avx);
}

static inline double
psnr(double d)
{
    if(d==0) return INFINITY;
    return log(d) / avx;
}

int
main(int argc, char *argv[])
{
    unsigned long plot = 0;
    unsigned long bufsize = 1<<20;
    Sample *samples = malloc(bufsize * sizeof(Sample));
    int frame, size;
    char t;
    double q, y, u, v, a;
    double lys = 0, lus = 0, lvs = 0, las = 0;
    double inf = 60.0;
    double lowcut = 0.0;
    double min = inf, max = 0;
    int count = 0;
    int lowct = 0;
    int lowsize = 0;
    int infct = 0;
    int infsize = 0;
    int less40 = 0;
    int max_frames = INT_MAX;
    unsigned long total_size = 0;
    int n;
    FILE *edl = 0;
    FILE *psnr_file = stdin;
    long start = -99, end = -99;
    unsigned char tt[256];
    double fps = 23.976;

    memset(tt, 1, sizeof(tt));
    avx = log(10)/-10.0;

    while (argc > 1 && argv[1][0] == '-') {
	int sw = argv[1][1];
	char *p = argv[1] + 2;
	if (!strcmp(argv[1], "-fps") && argv[2]) {
	    fps = strtod(argv[2], 0);
	    argv++;
	    argc--;
	} else if (!strcmp(argv[1], "-edl") && argv[2]) {
	    edl = fopen(argv[2], "r");
	    if (!edl) {
		fprintf(stderr, "Unable to open edl file %s: %s\n",
			argv[2], strerror(errno));
		exit(1);
	    }
	    argv++;
	    argc--;
	} else if (sw == 't') {
	    memset(tt, 0, sizeof(tt));
	    while (*p)
		tt[(unsigned char)*p++] = 1;
	} else if (sw == 'p') {
	    plot = strtoul(p, 0, 0);
	    if (!plot || plot > 4096)
		plot = 100;
	} else if (sw == 'h') {
	    fputs(usage, stdout);
	    return 0;
	} else if (sw == 'x' || sw == 'i' || sw == 'l') {
	    double x;
	    if (!*p) {
		p = argv[2];
		argv++;
		argc--;
	    }
	    if (!p || !(x = strtod(p, 0))) {
		fprintf(stderr, "Bax exponent\n");
		return 1;
	    }
	    if (sw == 'x')
		avx *= x;
	    else if (sw == 'i')
		inf = x;
	    else
		lowcut = x;
	} else if (!(max_frames = strtol(p-1, 0, 0))) {
	    fprintf(stderr, "Unknown option -%c.  Use -h for help\n", sw);
	    return 1;
	}
	argv++;
	argc--;
    }

    if (argc > 1) {
	psnr_file = fopen(argv[1], "r");
	if (!psnr_file) {
	    fprintf(stderr, "Unable to open %s: %s\n",
		    argv[1], strerror(errno));
	    exit(1);
	}
    }

    do {
	double iy, iu, iv;
	char buf[1024];
	char *p = fgets(buf, sizeof(buf), psnr_file);
	if (!p)
	    break;
	if ((n = sscanf(buf, "%d, %lf, %d, %lf, %lf, %lf, %lf %c\n",
			&frame,&q,&size,&y, &u, &v,  &a,  &t)) != 8) {
	    /* Try to match x264 log=3 output */
	    int xq;
	    if (!(p = strstr(buf, "x264")) ||
		(n = sscanf(p, "%*[^f]frame=%d QP=%d %*[^S]Slice:"
			    "%c %*[^s]size=%d %*[^P]PSNR Y:%lf U:%lf V:%lf",
			    &frame, &xq, &t, &size, &y, &u, &v)) != 7)
		continue;
	    a = 0;
	}
	max_frames--;
	if (frame > end + 20 && edl) {
	    double s, e;
	    int m;
	    p = fgets(buf, sizeof(buf), edl);
	    n = p ? sscanf(buf, "%lf %lf %d", &s, &e, &m) : EOF;
	    if (n != 3) {
		if (n != EOF)
		    fprintf(stderr, "Illegal edl file\n");
		fclose(edl);
		edl = 0;
	    } else {
		start = s * 23.976 + 0.5;
		end   = e * 23.976 + 1.5;
	    }
	}
	if ((start <= frame && frame < end) || !tt[(unsigned char)t])
	    continue;
	if (y < lowcut)
	    lowct++, lowsize += size;
	else if (y > inf)
	    infct++, infsize += size;
	else {
	    if (y < 40.0)
		less40++;
	    if (u > inf)
		u = inf;
	    if (v > inf)
		v = inf;
	    if (a > inf)
		a = inf;
	    lys += iy = invpsnr(y);
	    lus += iu = invpsnr(u);
	    lvs += iv = invpsnr(v);
	    las += a ? invpsnr(a) : (iy + iu/4 + iv/4) * (2.0/3);
	    if (y > max)
		max = y;
	    if (y < min)
		min = y;
	    if (count > bufsize) {
		bufsize <<= 1;
		samples = realloc(samples, bufsize * sizeof(Sample));
	    }
	    samples[count].r = y;
	    samples[count].size = size;
	    count++;
	}
	total_size += size;
    } while (max_frames);

    if (count) {
	printf("Y: %g, U: %g, V: %g, A: %g, %d+%d+%d = %d frames\n"
	       "size = %lu, infsize = %d; %.3f/%.3f kbps; %.2g%% < 40\n",
	       psnr(lys/count),
	       psnr(lus/count),
	       psnr(lvs/count),
	       psnr(las/count), lowct, count, infct, lowct+count+infct,
	       total_size, infsize,
	       total_size * (.008 * fps) / (infct + count),
	       (total_size - infsize) * (.008 * fps) / count,
	       100.0 * less40 / count);
    }

    if (plot) {
	unsigned (*counts)[2] = calloc(plot, 2 * sizeof(unsigned));
	int pp[2];
	FILE *cmd;
	int child;
	double scale;
	if (pipe(pp)) {
	    fprintf(stderr, "pipe() failed: %s\n", strerror(errno));
	    return 1;
	}
	switch ((child = fork())) {
	  case -1:
	    close(pp[0]);
	    close(pp[1]);
	    fprintf(stderr, "fork() failed: %s\n", strerror(errno));
	    return 1;
	  case 0:
	    close(0);
	    dup(pp[0]);
	    close(pp[0]);
	    close(pp[1]);
	    execlp("gnuplot", "gnuplot", NULL);
	    fprintf(stderr, "execlp() failed: %s\n", strerror(errno));
	    return 0;
	}
	cmd = fdopen(pp[1], "w");
	if (!cmd) {
	    fprintf(stderr, "fdopen() failed: %s\n", strerror(errno));
	    kill(child, SIGKILL);
	    return 1;
	}
	scale = (plot - 0.001) / (max - min);
	for (n = 0; n != count; n++) {
	    int i = (samples[n].r - min) * scale;
	    counts[i][0]++;
	    counts[i][1] += samples[n].size;
	}
	fprintf(cmd, "plot '-' smooth cspline title \"time\", "
		"'-' smooth cspline title \"size\"\n");
	total_size -= infsize;
	for (n = 0; n != plot; n++) {
	    fprintf(cmd, "%g %g\n", min + (n+.5)/scale,
		    scale * counts[n][0] / count);
	}
	fprintf(cmd, "e\n");
	for (n = 0; n != plot; n++) {
/*  	    if (counts[n][0]) */
	    fprintf(cmd, "%g %g\n", min + (n+.5)/scale,
		    scale * counts[n][1] / total_size
/*  		    !counts[n][0] ? .1 : */
/*  		    .1 * counts[n][1] * count / total_size / counts[n][0] */
		);
	}
	fprintf(cmd, "e\n");
	fflush(cmd);
	{
	    FILE *tty = fopen("/dev/tty", "r");
	    char junk[1024];
	    fgets(junk, 1024, tty);
	}
    }
    return 0;
}