//#include <afxwin.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include "bayer.h"

#define DEBUG_SNR 0

void CreateUV( unsigned char *Frame, int Width, int Height );
void PreProcUV(unsigned char *Frame, int Width, int Height);
void Green6( unsigned char *Frame, int Width, int Height );
void Demosaicing( unsigned char *Frame, int Width, int Height, FILE *fp );

//! Global Variables
int TAP9=0; //Disabled filter
void tvg_line_ext2( unsigned char *l, int Width )//by Fan Zhang
{
	//symmetric extention
	l[-1] = l[1];
	l[-2] = l[2];
	l[Width] = l[Width-2];
	l[Width+1] = l[Width-3];
}

int *U=NULL,*V=NULL,*Uh=NULL,*Vh=NULL,*Uv=NULL,*Vv=NULL;

//! new algorithm
void Green6( unsigned char *Frame, int Width, int Height )
{
	U = (int*) calloc(Width*Height,sizeof(int));
	V = (int*) calloc(Width*Height,sizeof(int));
	CreateUV(Frame,Width,Height);
}

typedef unsigned char pixel; 
void CreateUV( pixel *Frame, int Width, int Height )
{
	//const double Factor = 1.0/6.0;
	const double Factor = 1.0/4.0;
	pixel *line0,*line1,*line2,*line3,*line4;
	pixel *cLine,*nLine,*sLine,*nnLine,*ssLine;
	pixel *pBuf;
	int i,j,pos,band;
	int n,nn,ww,w,c,e,ee,s,ss;
	
	line0 = (pixel *) calloc(Width+4,sizeof(pixel));
	line1 = (pixel *) calloc(Width+4,sizeof(pixel));
	line2 = (pixel *) calloc(Width+4,sizeof(pixel));
	line3 = (pixel *) calloc(Width+4,sizeof(pixel));	
	line4 = (pixel *) calloc(Width+4,sizeof(pixel));
	
	Uh = (int *)calloc( Width*Height,sizeof(int)); 
	Vh = (int *)calloc( Width*Height,sizeof(int)); 
	Uv = (int *)calloc( Width*Height,sizeof(int)); 
	Vv = (int *)calloc( Width*Height,sizeof(int)); 
	//! HV prediction
	cLine=line0+2; nLine=line1+2; sLine=line2+2;
	nnLine=line3+2; ssLine=line4+2;
	pBuf = Frame;
	for( i=0; i<Height; i++ ) {	//symmetric expand
		memcpy(cLine,pBuf,Width*sizeof(pixel));
		if( i-1<0 )
			memcpy(nLine,pBuf+Width,Width*sizeof(pixel));
		else 
		    memcpy(nLine,pBuf-Width,Width*sizeof(pixel));
		if( i+1>=Height )
			memcpy(sLine,pBuf-Width,Width*sizeof(pixel));
		else
			memcpy(sLine,pBuf+Width,Width*sizeof(pixel));
		if( i-2<0 )
			memcpy(nnLine,pBuf+2*Width,Width*sizeof(pixel));
		else
			memcpy(nnLine,pBuf-2*Width,Width*sizeof(pixel));
		if( i+2>=Height )
			memcpy(ssLine,pBuf-2*Width,Width*sizeof(pixel));
		else
			memcpy(ssLine,pBuf+2*Width,Width*sizeof(pixel));
		tvg_line_ext2(cLine,Width); 
		tvg_line_ext2(nLine,Width); 
		tvg_line_ext2(sLine,Width); 
		tvg_line_ext2(nnLine,Width);
		tvg_line_ext2(ssLine,Width);
		//! GRGRGRGRGR 
		//! BGBGBGBGBG
		for( j=0; j<Width; j++ ) {
			double t;
			pos = i*Width+j;
			band = ((i&1)<<1)+(j&1);
			ww = cLine[j-2]; w = cLine[j-1]; 
			c = cLine[j];  e = cLine[j+1]; ee = cLine[j+2];
			nn = nnLine[j]; n = nLine[j];  
			s = sLine[j];  ss = ssLine[j];  
			if( band==BAYER_R ) {
				t = c-((c+c-ww-ee)*Factor)-((w+e)>>1); 
				Uh[pos] = (int)t;//R-G
				t = c-((c+c-nn-ss)*Factor)-((n+s)>>1); 
				Uv[pos] = (int)t;//R-G
			}
			if( band==BAYER_B ) {
				t = c-((c+c-ww-ee)*Factor)-((w+e)>>1); 
				Vh[pos] = (int)t;//B-G
				t = c-((c+c-nn-ss)*Factor)-((n+s)>>1); 
				Vv[pos] = (int)t;//B-G
			}
			if( band==BAYER_GR ) {//与paper上不符
				t = -c+((c+c-ww-ee)*Factor)+((w+e)>>1); 
				Uh[pos] = (int)t;//R-G
				t = -c+((c+c-nn-ss)*Factor)+((n+s)>>1); 
				Vv[pos] = (int)t;//B-G
			}
			if( band==BAYER_GB ) {//与paper上不符
				t = -c+((c+c-ww-ee)*Factor)+((w+e)>>1); 
				Vh[pos] = (int)t;//B-G
				t = -c+((c+c-nn-ss)*Factor)+((n+s)>>1); 
				Uv[pos] = (int)t;//R-G
			}
		}
		pBuf+=Width;
	}
	for( i=0; i<Height; i++ ) {
		for( j=0; j<Width; j++ ) {
			int pos = i*Width+j;
			band = ((i&1)<<1)+(j&1);
			if( i<1 || i+1>=Height || j<1 || j+1>=Width ) continue;//不处理边界
			if( band==BAYER_GR ) {
				Uv[pos] = (Uv[pos-1]+Uv[pos+1])>>1;
				Vh[pos] = (Vh[pos-Width]+Vh[pos+Width])>>1;
			}
			if( band==BAYER_R ) {
				Vh[pos] = (Vh[pos-1-Width]+Vh[pos+1-Width]+Vh[pos-1+Width]+Vh[pos+1+Width])>>2;
				Vv[pos] = (Vv[pos-1-Width]+Vv[pos+1-Width]+Vv[pos-1+Width]+Vv[pos+1+Width])>>2;
			}
			if( band==BAYER_B ) {
				Uh[pos] = (Uh[pos-1-Width]+Uh[pos+1-Width]+Uh[pos-1+Width]+Uh[pos+1+Width])>>2;
				Uv[pos] = (Uv[pos-1-Width]+Uv[pos+1-Width]+Uv[pos-1+Width]+Uv[pos+1+Width])>>2;
			}
			if( band==BAYER_GB ) {
				Vv[pos] = (Vv[pos-1]+Vv[pos+1])>>1;
				Uh[pos] = (Uh[pos-Width]+Uh[pos+Width])>>1;
			}
		}
	}
	if( TAP9 )
		PreProcUV(Frame, Width, Height);
	free(line0); free(line1); free(line2); free(line3); free(line4);
}

int LFilter( int *v, int taps )
{
	double t;
	t = 0;
	assert( taps==5 );
	t = (v[0]+v[1]+v[2]+v[3]+v[4])/5.0;
	if(t<0) t -= 0.5; 
	else t += 0.5;
	return((int)t);
}

void PreProcUV(unsigned char *F, int Width, int Height)//Directional Filter
{
	int i,j,k;
	int pos,band;
	int *uh,*uv,*vh,*vv;
	uh = (int *)calloc( Width*Height,sizeof(int)); 
	uv = (int *)calloc( Width*Height,sizeof(int)); 
	vh = (int *)calloc( Width*Height,sizeof(int)); 
	vv = (int *)calloc( Width*Height,sizeof(int)); 
	memcpy(uh,Uh,Width*Height*sizeof(int));
	memcpy(uv,Uv,Width*Height*sizeof(int));
	memcpy(vh,Vh,Width*Height*sizeof(int));
	memcpy(vv,Vv,Width*Height*sizeof(int));
	int h0[7],h1[7],v0[7],v1[7];
	for( i=0; i<Height; i++ ) {
		for( j=0; j<Width; j++ ) {
			pos = i*Width+j;
			band = ((i&1)<<1)+(j&1);
			if( i<3 || i+3>=Height || j<3 || j+3>=Width ) continue;
			for( k=-2; k<=2; k++ ) {
				h0[k+2] = Uh[pos+k];
				h1[k+2] = Vh[pos+k];
				v0[k+2] = Uv[pos+k*Width];
				v1[k+2] = Vv[pos+k*Width];
			}
			if( band==BAYER_R ) {
				uh[pos] = LFilter(h0,5);
				uv[pos] = LFilter(v0,5);
			}
			if( band==BAYER_B ) {
				vh[pos] = LFilter(h1,5);
				vv[pos] = LFilter(v1,5);
			}
		}
	}
	memcpy(Uh,uh,Width*Height*sizeof(int));
	memcpy(Uv,uv,Width*Height*sizeof(int));
	memcpy(Vh,vh,Width*Height*sizeof(int));
	memcpy(Vv,vv,Width*Height*sizeof(int));
	free(uh); free(uv); free(vh); free(vv);
}

int grad_h( int h[9] )
{
	int grad = 0;
	//0 1 2
	//3 4 5
	//6 7 8
	grad = abs(h[0]-h[1])+abs(h[1]-h[2])+abs(h[0]-h[2]);
	grad+= abs(h[3]-h[4])+abs(h[4]-h[5])+abs(h[3]-h[5]);
	grad+= abs(h[6]-h[7])+abs(h[7]-h[8])+abs(h[6]-h[8]);
/*
	grad+= abs(h[0]-h[3])+abs(h[6]-h[3])
		  +abs(h[1]-h[4])+abs(h[7]-h[4])
		  +abs(h[2]-h[5])+abs(h[8]-h[5]);
*/
	return(grad);
}

int grad_v( int v[9] )
{
	int grad = 0;
	//0 1 2
	//3 4 5
	//6 7 8
	grad = abs(v[0]-v[3])+abs(v[3]-v[6])+abs(v[0]-v[6]);
	grad+= abs(v[1]-v[4])+abs(v[4]-v[7])+abs(v[1]-v[7]);
	grad+= abs(v[2]-v[5])+abs(v[5]-v[8])+abs(v[2]-v[8]);
/*
	grad+= abs(v[0]-v[1])+abs(v[2]-v[1])
		  +abs(v[3]-v[4])+abs(v[5]-v[4])
		  +abs(v[6]-v[7])+abs(v[8]-v[7]);
*/
	return(grad);
}

void IPT_Demosaicing( unsigned char *Frame, int Width, int Height, unsigned char *out_buf )
{
	#define W Width
	#define F Frame
	unsigned char *rgbLine;
	unsigned char *p_out_buf = out_buf;
	int i,j;
	int uh[9],uv[9],vh[9],vv[9];
#if DEBUG_SNR
	extern unsigned char *rIntra,*gIntra,*bIntra;
	assert(rIntra!=NULL && gIntra!=NULL && bIntra!=NULL);
#endif		
	//! do all 6 interpolations
	Green6( Frame, Width, Height );
	//! HV decision
	for( i=0; i<Height; i++ ) {
		//if((i%100)==0) printf("%d ",i);
		for(j=0;j<Width;j++){
			int pos = i*Width+j; 
			int band = ((i&1)<<1)+(j&1);
			if( i>=2 && i+2<Height && j>=2 && j+2<Width ) {
				double hAve,vAve,hVar,vVar;
				int h[5],v[5],k;
				int *Ch,*Cv,*dstC;
				if( band==BAYER_GR || band==BAYER_GB ) continue;
				int w,e,c,ww,ee,n,nn,s,ss;
uh[0] = Uh[pos-W-1]; uh[1] = Uh[pos-W]; uh[2] = Uh[pos-W+1];
uh[3] = Uh[pos-1  ]; uh[4] = Uh[  pos]; uh[5] = Uh[pos+1];
uh[6] = Uh[pos+W-1]; uh[7] = Uh[pos+W]; uh[8] = Uh[pos+W+1];
uv[0] = Uv[pos-W-1]; uv[1] = Uv[pos-W]; uv[2] = Uv[pos-W+1];
uv[3] = Uv[pos-1  ]; uv[4] = Uv[  pos]; uv[5] = Uv[  pos+1];
uv[6] = Uv[pos+W-1]; uv[7] = Uv[pos+W]; uv[8] = Uv[pos+W+1];
vh[0] = Vh[pos-W-1]; vh[1] = Vh[pos-W]; vh[2] = Vh[pos-W+1];
vh[3] = Vh[pos-1  ]; vh[4] = Vh[  pos]; vh[5] = Vh[pos+1  ];
vh[6] = Vh[pos+W-1]; vh[7] = Vh[pos+W]; vh[8] = Vh[pos+W+1];
vv[0] = Vv[pos-W-1]; vv[1] = Vv[pos-W]; vv[2] = Vv[pos-W+1];
vv[3] = Vv[pos-1  ]; vv[4] = Vv[pos  ]; vv[5] = Vv[pos+1];
vv[6] = Vv[pos+W-1]; vv[7] = Vv[pos+W]; vv[8] = Vv[pos+W+1];
				w = Frame[pos-1]; c = Frame[pos]; e = Frame[pos+1]; 
				ww = Frame[pos-2]; ee = Frame[pos+2]; 
				n = Frame[pos-Width]; s = Frame[pos+Width]; 
				nn = Frame[pos-2*Width]; ss = Frame[pos+2*Width]; 
				int dh = abs(w-e)+abs(c+c-ww-ee);
				int dv = abs(n-s)+abs(c+c-nn-ss);
dh = grad_h(uh) + grad_h(vh);
dv = grad_v(uv) + grad_v(vv);
				
				if( band==BAYER_R ) {
					Ch = Uh; Cv = Uv; dstC=U;
				}
				if( band==BAYER_B ) {
					Ch = Vh; Cv = Vv; dstC=V;
				}
				
				for( k=-2; k<=2; k++ ) {
					h[k+2] = Ch[pos+k];
				}
				for( k=-2; k<=2; k++ ) {
					v[k+2] = Cv[pos+k*Width];
				}
				hAve = (h[0]+h[1]+h[2]+h[3]+h[4])/5.0;
				vAve = (v[0]+v[1]+v[2]+v[3]+v[4])/5.0;
				hVar=vVar=0;
				for( k=0; k<5; k++ ) {
					hVar += (h[k]-hAve)*(h[k]-hAve);
					vVar += (v[k]-vAve)*(v[k]-vAve);
				}
				int ph,pv;
				ph = Ch[pos]; pv = Cv[pos];
				double pw;
				//ph = pv;
#if 0				
				if( dh<dv ) {
					pv = ph;
				}
				else if ( dv<dh ) {
					ph = pv;
				}
				else {
					int t = (ph+pv)>>1;
					ph=pv=t;
				}
#endif
				if( hVar==0 ) {
					pw = ph;
				}
				else if( vVar==0 ) {
					pw = pv;
				}
				else 
					pw = (ph*vVar+pv*hVar)/(hVar+vVar);
				if( pw<F[pos]-255 ) pw=F[pos]-255; 
				if( pw>F[pos] ) pw = F[pos];
				if( pw<0 ) pw -= 0.5;
				if( pw>0 ) pw += 0.5;
				dstC[pos] = (int)pw;
//				Ch[pos] = Cv[pos] = (int)pw;
			}
		}
	}
	for( i=0; i<Height; i++ ) {
		//! GR..BG
		int band;
		for( j=0; j<Width; j++ ) {
			int pos = i*Width+j;
			band = ((i&1)<<1)+(j&1);
			if( i>=3 && i+3<Height && j>=3 && j+3<Width ) {
				if( band==BAYER_GR ) {
					U[pos] = (U[pos-1]+U[pos+1])>>1;
					V[pos] = (V[pos-W]+V[pos+W])>>1;
				}
				if( band==BAYER_R ) {
					V[pos] = (V[pos-W-1]+V[pos-W+1]+V[pos+W-1]+V[pos+W+1])>>2;
                }
				if( band==BAYER_B ) {
					U[pos] = (U[pos-W-1]+U[pos-W+1]+U[pos+W-1]+U[pos+W+1])>>2;
				}
				if( band==BAYER_GB ) {
					V[pos] = ((V[pos-1]+V[pos+1])>>1);
					U[pos] = ((U[pos-W]+U[pos+W])>>1);
				}
			}
		}
	}
	rgbLine = (unsigned char *) calloc(3*Width,sizeof(unsigned char));
	for( i=0; i<Height; i++ ) {
		//! GR..BG
		int r,g,b,band;
		for( j=0; j<Width; j++ ) {
			int pos = i*Width+j;
			band = ((i&1)<<1)+(j&1);
			unsigned char uct=0;
			//! special cases
			if( i>=3 && i+3<Height && j>=3 && j+3<Width )//why??
			{
				if( band==BAYER_GR ) {
					g = F[pos];  r = g+U[pos];  b = g+V[pos];
				}
				if( band==BAYER_R ) {
					r = F[pos];  g = F[pos] - U[pos];  b = g + V[pos];
                }
				if( band==BAYER_B ) {
					g = F[pos] - V[pos];  b = F[pos];  r = g+U[pos];
				}
				if( band==BAYER_GB ) {
					g = F[pos];  b = g+V[pos];  r = g+U[pos];
				}
				if( r<0 ) r=0; else if ( r>255) r=255;
				if( g<0 ) g=0; else if ( g>255) g=255;
				if( b<0 ) b=0; else if ( b>255) b=255;
			}
			else {
				r=g=b=0;
			}
#if DEBUG_SNR
			rIntra[pos]=r; gIntra[pos]=g; bIntra[pos]=b;
#endif			
			rgbLine[3*j+0]=b; rgbLine[3*j+1]=g; rgbLine[3*j+2]=r; 
		}
		memcpy(p_out_buf, rgbLine, 3*Width*sizeof(unsigned char) );
		p_out_buf += 3*Width;
	}
	free(rgbLine); 
	free(Uh); free(Uv); free(Vh); free(Vv); free(U); free(V); 
}

