#include <stdio.h>
#include <math.h>
#include "config.h"
#include "global.h"

/* private prototypes */
static void calcSNR1 _ANSI_ARGS_((unsigned char *org, unsigned char *rec,
  int lx, int w, int h, double *pv, double *pe));


void calcSNR(org,rec)
unsigned char *org[3];
unsigned char *rec[3];
{
  int w,h,offs;
  double v,e;

  w = horizontal_size;
  h = (pict_struct==FRAME_PICTURE) ? vertical_size : (vertical_size>>1);
  offs = (pict_struct==BOTTOM_FIELD) ? width : 0;

  calcSNR1(org[0]+offs,rec[0]+offs,width2,w,h,&v,&e);
  fprintf(statfile,"Y: variance=%4.4g, MSE=%3.3g (%3.H3g dB), SNR=%3.3g dB\n",
    v, e, 10.0*log10(255.0*255.0/e), 10.0*log10(v/e));
    
    printf("Y: variance=%4.4g, MSE=%3.3g (%3.H3g dB), SNR=%3.3g dB\n", v, e, 10.0*log10(255.0*255.0/e), 10.0*log10(v/e));
    
    

  if (chroma_format!=CHROMA444)
  {
    w >>= 1;
    offs >>= 1;
  }

  if (chroma_format==CHROMA420)
    h >>= 1;

  calcSNR1(org[1]+offs,rec[1]+offs,chrom_width2,w,h,&v,&e);
  fprintf(statfile,"U: variance=%4.4g, MSE=%3.3g (%3.3g dB), SNR=%3.3g dB\n",
    v, e, 10.0*log10(255.0*255.0/e), 10.0*log10(v/e));
  printf("U: variance=%4.4g, MSEFFF=%3.3g (%3.3g dB), SNR=%3.3g dB\n", v, e, 10.0*log10(255.0*255.0/e), 10.0*log10(v/e));
    

  calcSNR1(org[2]+offs,rec[2]+offs,chrom_width2,w,h,&v,&e);
  fprintf(statfile,"V: variance=%4.4g, MSE=%3.3g (%3.3g dB), SNR=%3.3g dB\n",
    v, e, 10.0*log10(255.0*255.0/e), 10.0*log10(v/e));
    
 printf("V: variance=%4.4g, MSE=%3.3g (%3.3g dB), SNR=%3.3g dB\n", v, e, 10.0*log10(255.0*255.0/e), 10.0*log10(v/e));    
}

static void calcSNR1(org,rec,lx,w,h,pv,pe)
unsigned char *org;
unsigned char *rec;
int lx,w,h;
double *pv,*pe;
{
  int i, j;
  double v1, s1, s2, e2;

  s1 = s2 = e2 = 0.0;

  for (j=0; j<h; j++)
  {
    for (i=0; i<w; i++)
    {
      v1 = org[i];
      s1+= v1;
      s2+= v1*v1;
      v1-= rec[i];
      e2+= v1*v1;
    }
    org += lx;
    rec += lx;
  }

  s1 /= w*h;
  s2 /= w*h;
  e2 /= w*h;

  /* prevent division by zero in calcSNR() */
  if(e2==0.0)
    e2 = 0.00001;

  *pv = s2 - s1*s1; /* variance */
  *pe = e2;         /* MSE */
}