#include <stdlib.h>
#include <stdio.h>
#include "highgui.h"
#include "cv.h"

#define SQUARE(x) ((x)*(x))

double CPSNR = 0;
double PSNR_R = 0;
double PSNR_G = 0;
double PSNR_B = 0;

void IPT_Demosaicing( unsigned char *Frame, int Width, int Height, unsigned char *out_buf );

void RGB24_to_BayerCFA(unsigned char* RGB ,unsigned char* Bayer_CFA, int w, int h)
{
	int i,j,k;
	unsigned char G1,G2,R,B;

	for( k =0; k < h; k+=2 )
	{
		for( i =0, j=0; i < w; i+=2, j+=6 )
		{
			G1 = RGB[3*w*k +j+1];
			G2 = RGB[3*w*(k+1) +j+4];
			R = RGB[3*w*k +j+3+2];//
			B = RGB[3*w*(k+1) +j];//

			Bayer_CFA[w*k +i] = G1;
			Bayer_CFA[w*k +i+1] = R;
			Bayer_CFA[w*(k+1) +i] = B;
			Bayer_CFA[w*(k+1) +i+1] = G2;
		}
	}
}

void PSNR_raw_Crop(unsigned char* data1_BGR ,unsigned char* data2_BGR, const int w, const int h)
{
	double MSE_R=0;
	double MSE_G=0;
	double MSE_B=0;
	double CMSE=0;

	int i,j,k;
	for( i =3; i < h-3; i++ )
	{
		for( j =3; j < w-3; j++ )
		{
			k = 3*(i*w + j);
			MSE_R += SQUARE((double)data1_BGR[k +2] - data2_BGR[k +2]);
			MSE_G += SQUARE((double)data1_BGR[k +1] - data2_BGR[k +1]);
			MSE_B += SQUARE((double)data1_BGR[k] - data2_BGR[k]);
		}
	}
	CMSE = (MSE_B + MSE_G + MSE_R) /3;

	PSNR_R += 10 * log10( 255*255 / (MSE_R/(w-6)/(h-6)) );
	PSNR_G += 10 * log10( 255*255 / (MSE_G/(w-6)/(h-6)) );
	PSNR_B += 10 * log10( 255*255 / (MSE_B/(w-6)/(h-6)) );
	CPSNR += 10 * log10( 255*255 / (CMSE/(w-6)/(h-6)) );
}

int main(int argc,char **argv) // I/O for testing
{
	if (argc != 3)
	{
		printf("wrong argument number.");
		return -1;
	}

	IplImage* origin_img = cvLoadImage( argv[1], CV_LOAD_IMAGE_COLOR );
	if(! origin_img)
	{	
		printf("Error in image reading.");
		return -1;
	}

	IplImage* out_image = cvCloneImage(origin_img);

	unsigned char* origin_buf = (unsigned char*)origin_img->imageData;
	unsigned char* out_buf = (unsigned char*)out_image->imageData;

	int width = origin_img->width;
	int height= origin_img->height;

	unsigned char* in_buf = (unsigned char*)malloc(height * width);

	RGB24_to_BayerCFA(origin_buf, in_buf, width, height); //Mosaic

	IPT_Demosaicing(in_buf, width, height, out_buf);

	cvSaveImage(argv[2], out_image);


	//Calculate PSNR
	PSNR_raw_Crop(origin_buf, out_buf, width, height);
	printf("%s\tPSNR_R = %.2f\tPSNR_G = %.2f\tPSNR_B = %.2f\tCPSNR = %.2f\r\n",argv[1],PSNR_R,PSNR_G,PSNR_B,CPSNR);

	cvReleaseImage(&origin_img);
	cvReleaseImage(&out_image);

	//system("PAUSE");
	return 0;
}