#define WIN32_LEAN_AND_MEAN
#include <stdio.h>
#include <windows.h>
#include <crtdbg.h>
#include "avisynth.h"
#include <string.h>
#include "VqmCalc.h"
#include "math.h"

VqmCalc::VqmCalc(PClip _child1, PClip _child2, const char *fname, IScriptEnvironment* env) :
  GenericVideoFilter(_child1), child2(_child2), log(NULL) {

  const VideoInfo& vi2 = child2->GetVideoInfo();

  notdone=true;
  VqmFrame= 0;

  blockA= new short[64];
  blockB= new short[64];
  blockA2= new float[64];
  blockB2= new float[64];
  fdct = fdct_mmx;

  if (fname[0] != 0) {
    log = fopen(fname, "wt");
    if (log && notdone) {
      fprintf(log,"VqmCalc: Video Quality Metric 0.1\n");
    } else
      env->ThrowError("VQMcalc: unable to create file %s", fname);
  }
  
  if (!vi.IsYV12()) 
    env->ThrowError("Compare: YV12 colorspace required");

 
  
  width = vi.width;
  widthUV = width/2;
  height = vi.height;
  heightUV = height/2;
  
}

VqmCalc::~VqmCalc() {
  if (log && notdone) {
    fprintf(log,"Average Vqm= %f\n",(VqmFrame/((float)vi.num_frames)));
    fclose(log);
    notdone=false;
  }
  delete [] blockA;
  delete [] blockB;
  delete [] blockA2;
  delete [] blockB2;
} 

void VqmCalc::FillBlocks(const BYTE* ref,short* outref,int pitchref,const BYTE* alt,short* outalt,int pitchalt) {
  const BYTE* y;
  const BYTE* z;
  int i,j;
  
  y=ref;
  z=alt;
  j=0;
  for(i=0;i<64;i++){
    if(j==8){
      j=0;y+=pitchref;z+=pitchalt;}
    outref[i]=(short) y[j];
    outalt[i]=(short) z[j];
    j++;
  }  
}

void VqmCalc::LocalContrasts(float* ref, float* alt) {
  float dcref,dcalt;
  int i;
  dcref= powf((ref[0]/1024),0.65f)/ref[0];
  dcalt= powf((alt[0]/1024),0.65f)/alt[0];
  for(i=0;i<64;i++) {
    ref[i] *= dcref;
    alt[i] *= dcalt;
  }
}

void VqmCalc::ApplySCSFMatrix(float* ref, float* alt) {
  int i;
  for(i=0;i<64;i++){
    ref[i]/=((float)mpegmatrix[i]);
    alt[i]/=((float)mpegmatrix[i]);
  }
}

void VqmCalc::DiffBlocks(float* ref, float* alt) {
  int i;
  for(i=0;i<64;i++)
    ref[i]=fabsf(ref[i]-alt[i]);
}

float VqmCalc::MeanBlock(float* block) {
  float z= block[0];
  int i;
  for(i=1;i<64;i++) 
    z+=block[i];
  return z/64.0f;  
}

float VqmCalc::MaxBlock(float* block) {
  float z= block[0];
  int i;
  for(i=1;i<64;i++) 
    if(z<block[i]) 
      z=block[i];
  return z;
}

void VqmCalc::ConvertTo(short* ref, float* fref, short* alt, float* falt){
  int i;
  for(i=0;i<64;i++){
    fref[i]=(float) ref[i];
    falt[i]=(float) alt[i];
  }
}

PVideoFrame __stdcall VqmCalc::GetFrame(int n, IScriptEnvironment* env) {
  
  PVideoFrame refP = child->GetFrame(n, env);
  PVideoFrame altP = child2->GetFrame(n, env);
  
  const BYTE* refY = refP->GetReadPtr();
  const BYTE* altY = altP->GetReadPtr();
  const BYTE* refU = refP->GetReadPtr(PLANAR_U);
  const BYTE* altU = altP->GetReadPtr(PLANAR_U);
  const BYTE* refV = refP->GetReadPtr(PLANAR_V);
  const BYTE* altV = altP->GetReadPtr(PLANAR_V);
  
  const int refPitch = refP->GetPitch();
  const int refPitchUV = refP->GetPitch(PLANAR_U);
  const int altPitch = altP->GetPitch();
  const int altPitchUV = altP->GetPitch(PLANAR_U);

  int i,j;
  float a,b,c,d;
  int xiu;

  if(log && notdone){    
    fprintf(log,"n=%d",n);
  }

  imagm=0;
  imagM=0;
 
  for(i=0;i<height;i+=8) {
    //fprintf(log,"\nheight=%d\n",i);
    for(j=0;j<width;j+=8) {
      //fprintf(log,"\nwidth=%d\n",j);
	FillBlocks(refY+j,blockA,refPitch,altY+j,blockB,altPitch);
	//for(xiu=0;xiu<64;xiu++)
	//fprintf(log," i[%d]=%d ",xiu,blockA[xiu]-blockB[xiu]);
	//fprintf(log,"\n");
	fdct(blockA);
	fdct(blockB);
	ConvertTo(blockA,blockA2,blockB,blockB2);
	LocalContrasts(blockA2,blockB2);
  	ApplySCSFMatrix(blockA2,blockB2);
 	DiffBlocks(blockA2,blockB2);
	imagm+=MeanBlock(blockA2);
	imagM=max(imagM,MaxBlock(blockA2));
	//fprintf(log,"\timagm=%f imagM=%f\n",imagm,imagM);
    }
    refY+=(8*refPitch);
    altY+=(8*altPitch);
  }
  
  a=10*((64000*imagm/((float)(height*width)))+5*imagM);
  //fprintf(log,"\timagm=%f imagM=%f\n",imagm,imagM);
  if(log && notdone) {
    fprintf(log,"\tY=%f",a);
  }
  
  imagm=0;
  imagM=0;
  for(i=0;i<heightUV;i+=8) {
    for(j=0;j<widthUV;j+=8) {
      FillBlocks(refU+j, blockA, refPitchUV, altU+j, blockB, altPitchUV);
      fdct(blockA);
      fdct(blockB);
      ConvertTo(blockA,blockA2,blockB,blockB2);
      LocalContrasts(blockA2,blockB2);
      ApplySCSFMatrix(blockA2,blockB2);
      DiffBlocks(blockA2,blockB2);
      imagm+=MeanBlock(blockA2);
      imagM=max(imagM,MaxBlock(blockA2));
    }
    refU+=(8*refPitchUV);
    altU+=(8*altPitchUV);
  }
  b=10*((64000*imagm/((float)(heightUV*widthUV)))+5*imagM);
  if(log && notdone) {
    fprintf(log,"\tU=%f",b);
  }
  
  imagm=0;
  imagM=0;
  for(i=0;i<heightUV;i+=8) {
    for(j=0;j<widthUV;j+=8) {
      FillBlocks(refV, blockA, refPitchUV, altV, blockB, altPitchUV);
      fdct(blockA);
      fdct(blockB);
      ConvertTo(blockA,blockA2,blockB,blockB2);
      LocalContrasts(blockA2,blockB2);
      ApplySCSFMatrix(blockA2,blockB2);
      DiffBlocks(blockA2,blockB2);
      imagm+=MeanBlock(blockA2);
      imagM=max(imagM,MaxBlock(blockA2));
    }
    refV+=(8*refPitchUV);
    altV+=(8*altPitchUV);
  }
  c=10*((64000*imagm/((float)(heightUV*widthUV)))+5*imagM);
  if(log && notdone) {
    fprintf(log,"\tV=%f\n",c);
  }
  
  d=((0.8f*a)+(0.1f*b)+(0.1f*c));
  VqmFrame+=d;
  if(log && notdone) {
    fprintf(log,"VQM=%f\n",d);
  }
  
  return altP;
} 

AVSValue __cdecl Create_VqmCalc(AVSValue args, void* user_data, IScriptEnvironment* env) {
  return new VqmCalc(args[0].AsClip(), args[1].AsClip(), args[2].AsString(""), env);  
}

extern "C" __declspec(dllexport) const char* __stdcall AvisynthPluginInit2(IScriptEnvironment* env) {
  env->AddFunction("VqmCalc", "ccs", Create_VqmCalc, 0);
  return "`VqmCalc' VqmCalc plugin";
}