#include "mex.h"
#include "Mates.h"
#include <stdio.h>
#include <math.h>
#include <float.h>
#include <string.h>
#include <stdlib.h>

/* 

In order to compile this MEX function, type the following at the MATLAB prompt:
mex SteeringKernelRegressionSparse.c MatesLap.c Debugging.c lapack.a blas.a libf2c.a

Usage:
[RecImg, GradX, GradY, DXX, DXY, DYY] = SteeringKernelRegressionSparse(NoisyImg, Resp, h, C, ksize);

Notes:
RecImg=Reconstructed image, Z
GradX=Gradient in the horizontal direction, dZ/dX
GradY=Gradient in the vertical direction, dZ/dY
DXX=Second derivative with respect to the horizontal direction, d^2 Z / d X^2
DXY=Second derivative with respect to both directions, d^2 Z / dX dY
DYY=Second derivative with respect to the vertical direction, d^2 Z / d Y^2
NoisyImg=Input (noisy) image
Resp=Probability that a given pixel is not impulse corrupted
h=Kernel regression global smoothing parameter
C=steering matrices map
ksize=Kernel size

Version optimized for sparse Resp matrices

*/


/* Extend a 2D matrix (mirror extension at the edges) */
void Extend(double *Dest,double *Orig,int N, int M,int radius);

void mexFunction(int nlhs, mxArray* plhs[],
                 int nrhs, const mxArray* prhs[])
{  
	double *Responsibilities,*NoisyImg,*RecImg,*GradX,*GradY,*C,*DXX,*DXY,*DYY;
	double *Xx,*XxT,*Xw,*XwT,*XxTXw,*Inv,*A,*AT,*XCoord,*YCoord,*W;
	double *NoisyImgExt,*ResponsibilitiesExt,*Window,*WindowResp;
	double *WindowC11,*WindowC12,*WindowC22,*WindowSqrtDetC;
	double *C11,*C12,*C22,*C11Ext,*C12Ext,*C22Ext,*SqrtDetC,*SqrtDetCExt,*MyResp;
	double h;
	int radius,ksize,M,N,NdxWin,Ndx,Cnt,NdxCol,NdxRow,NonZeros,NdxNonZero,NdxWinCol,NdxWinRow;

    /* Get input data */
    h=mxGetScalar(prhs[2]);    
	ksize=(int)mxGetScalar(prhs[4]);    
	NoisyImg=mxGetPr(prhs[0]);
	Responsibilities=mxGetPr(prhs[1]);
	C=mxGetPr(prhs[3]);
    
    
    /* Create output arrays */
	N=mxGetM(prhs[0]);  /* Number of rows in the input image */
	M=mxGetN(prhs[0]);  /* Number of columns in the input image */
    plhs[0]=mxCreateDoubleMatrix(N,M,mxREAL);
	RecImg=mxGetPr(plhs[0]);
	plhs[1]=mxCreateDoubleMatrix(N,M,mxREAL);
	GradX=mxGetPr(plhs[1]);
	plhs[2]=mxCreateDoubleMatrix(N,M,mxREAL);
	GradY=mxGetPr(plhs[2]);
	plhs[3]=mxCreateDoubleMatrix(N,M,mxREAL);
	DXX=mxGetPr(plhs[3]);
	plhs[4]=mxCreateDoubleMatrix(N,M,mxREAL);
	DXY=mxGetPr(plhs[4]);
	plhs[5]=mxCreateDoubleMatrix(N,M,mxREAL);
	DYY=mxGetPr(plhs[5]);
    
    
    /* Get working variables */
	radius = (ksize-1)/2;
         
    /* Prepare auxiliary arrays */
    Xx=mxMalloc((ksize*ksize)*6*sizeof(double));
	XxT=mxMalloc(6*(ksize*ksize)*sizeof(double));
	Xw=mxMalloc((ksize*ksize)*6 *sizeof(double));
	XwT=mxMalloc(6*(ksize*ksize)*sizeof(double));
	XxTXw=mxMalloc(6*6*sizeof(double));
	Inv=mxMalloc(6*6*sizeof(double));
	A=mxMalloc(6*(ksize*ksize)*sizeof(double));
	AT=mxMalloc((ksize*ksize)*6*sizeof(double));
	XCoord=mxMalloc(ksize*ksize*sizeof(double));
	YCoord=mxMalloc(ksize*ksize*sizeof(double));	
	W=mxMalloc(ksize*ksize*sizeof(double));
	NoisyImgExt=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));
	ResponsibilitiesExt=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));
	C11=mxMalloc(N*M*sizeof(double));
	C12=mxMalloc(N*M*sizeof(double));
	C22=mxMalloc(N*M*sizeof(double));
	SqrtDetC=mxMalloc(N*M*sizeof(double));
	C11Ext=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));	
	C12Ext=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));	
	C22Ext=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));	
	SqrtDetCExt=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));	
	Window=mxMalloc(ksize*ksize*sizeof(double));
	WindowResp=mxMalloc(ksize*ksize*sizeof(double));
	WindowC11=mxMalloc(ksize*ksize*sizeof(double));
	WindowC12=mxMalloc(ksize*ksize*sizeof(double));
	WindowC22=mxMalloc(ksize*ksize*sizeof(double));
	WindowSqrtDetC=mxMalloc(ksize*ksize*sizeof(double));
	MyResp=mxMalloc(ksize*ksize*sizeof(double));


	/* Prepare window coordinates */
	for(NdxWin=0;NdxWin<ksize*ksize;NdxWin++)
	{
		XCoord[NdxWin]=(NdxWin%ksize)-radius;
		YCoord[NdxWin]=(NdxWin/ksize)-radius;
	}

	

	/* Prepare feature matrix */
	for(NdxWin=0;NdxWin<ksize*ksize;NdxWin++)
	{
		Xx[NdxWin]=1.0;
		Xx[ksize*ksize+NdxWin]=XCoord[NdxWin];
		Xx[2*ksize*ksize+NdxWin]=YCoord[NdxWin];
		Xx[3*ksize*ksize+NdxWin]=XCoord[NdxWin]*XCoord[NdxWin];
		Xx[4*ksize*ksize+NdxWin]=XCoord[NdxWin]*YCoord[NdxWin];
		Xx[5*ksize*ksize+NdxWin]=YCoord[NdxWin]*YCoord[NdxWin];
	}	

	/* Prepare steering matrices */
	for(Ndx=0;Ndx<N*M;Ndx++)
	{
		C11[Ndx]=C[4*Ndx];
		C12[Ndx]=C[4*Ndx+1];
		C22[Ndx]=C[4*Ndx+3];
		SqrtDetC[Ndx]=sqrt(C[4*Ndx]*C[4*Ndx+3]-C[4*Ndx+1]*C[4*Ndx+1]);
	}

	/* Extend the noisy image, the responsibilities map and the maps for steering matrices*/
	Extend(NoisyImgExt,NoisyImg,N,M,radius);	
	Extend(ResponsibilitiesExt,Responsibilities,N,M,radius);	
	Extend(C11Ext,C11,N,M,radius);
	Extend(C12Ext,C12,N,M,radius);
	Extend(C22Ext,C22,N,M,radius);
	Extend(SqrtDetCExt,SqrtDetC,N,M,radius);

	/*---------*/

	/* Estimate the original image and the gradients */
	for(NdxCol=0;NdxCol<M;NdxCol++)
	{
		for(NdxRow=0;NdxRow<N;NdxRow++)
		{
			/* Get the responsibilities in MyResp */
			for(NdxWin=0;NdxWin<ksize;NdxWin++)
			{
				memcpy(MyResp+NdxWin*ksize,ResponsibilitiesExt+(NdxCol+NdxWin)*(N+2*radius)+NdxRow,
					ksize*sizeof(double));
			}
			
			/* Count the number of pixels with nonzero responsibility */
			NonZeros=0; 
			for(NdxWin=0;NdxWin<ksize*ksize;NdxWin++)
			{
				if (MyResp[NdxWin]!=0.0)
				{
					NonZeros++;
				}
			}

			/* Get the relevant noisy samples and the rest of parameters,
			 restricted to pixels with nonzero responsibilities*/
			NdxNonZero=0;
			for(NdxWin=0;NdxWin<ksize*ksize;NdxWin++)
			{
				/* Process only pixels with nonzero responsibility */
				if (MyResp[NdxWin]!=0.0)				
				{
					/* Coordinates inside the window */
					NdxWinCol=NdxWin/ksize;
					NdxWinRow=NdxWin%ksize;

					/* Get the parameters */
					Window[NdxNonZero]=NoisyImgExt[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];
					WindowC11[NdxNonZero]=C11Ext[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];
					WindowC12[NdxNonZero]=C12Ext[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];
					WindowC22[NdxNonZero]=C22Ext[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];
					WindowSqrtDetC[NdxNonZero]=SqrtDetCExt[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];

					/* The responsibility */
					WindowResp[NdxNonZero]=MyResp[NdxWin];

					/* Traspose of the feature matrix */			
					for(Cnt=0;Cnt<6;Cnt++)
					{						
						XxT[6*NdxNonZero+Cnt]=Xx[ksize*ksize*Cnt+NdxWin];						
					}	

					/* Prepare weight */	
					W[NdxNonZero]=XxT[6*NdxNonZero+1]*
						(WindowC11[NdxNonZero]*XxT[6*NdxNonZero+1]+WindowC12[NdxNonZero]*XxT[6*NdxNonZero+2])
						+XxT[6*NdxNonZero+2]*
						(WindowC12[NdxNonZero]*XxT[6*NdxNonZero+1]+WindowC22[NdxNonZero]*XxT[6*NdxNonZero+2]);
					W[NdxNonZero]=exp(-(0.5/(h*h)) * W[NdxNonZero])*WindowSqrtDetC[NdxNonZero];

					/* The weight multiplied by the responsibility */
					WindowResp[NdxNonZero]=MyResp[NdxWin]*W[NdxNonZero];
					
					/* Equivalent kernel */			
					for(Cnt=0;Cnt<6;Cnt++)
					{
						Xw[NonZeros*Cnt+NdxNonZero]=Xx[ksize*ksize*Cnt+NdxWin]*WindowResp[NdxNonZero];							
					}	

					NdxNonZero++;
				}
			}

			/* Traspose the equivalent kernel */		
			Traspose(Xw,XwT,NonZeros,6);

			/* XxTXw = Xx^T * Xw */
			MatrixProduct(XxT,Xw,XxTXw,6,NonZeros,6);
			/* Add 0.00001 to the elements of the diagonal */
			SumDiagonalConstant(XxTXw,0.00001,NULL,6);
			/* Invert it */
			Inverse(XxTXw,Inv,6);

			/* A=Inv*(Xw^T)*/
			MatrixProduct(Inv,XwT,A,6,6,NonZeros);

			Traspose(A,AT,6,NonZeros);

			/* Estimate the original image and the gradients at this position */
			(*RecImg)=0.0;
			(*GradX)=0.0;
			(*GradY)=0.0;
			(*DXX)=0.0;
			(*DXY)=0.0;
			(*DYY)=0.0;
			for(NdxWin=0;NdxWin<NonZeros;NdxWin++)
			{
				(*RecImg)+=AT[NdxWin]*Window[NdxWin];
				(*GradX)+=AT[NonZeros+NdxWin]*Window[NdxWin];
				(*GradY)+=AT[2*NonZeros+NdxWin]*Window[NdxWin];
				(*DXX)+=AT[3*NonZeros+NdxWin]*Window[NdxWin];
				(*DXY)+=AT[4*NonZeros+NdxWin]*Window[NdxWin];
				(*DYY)+=AT[5*NonZeros+NdxWin]*Window[NdxWin];
			}
						
			/* Advance output pointers */
			RecImg++;
			GradX++;
			GradY++;
			DXX++;
			DXY++;
			DYY++;
		}

	}

	/* Release memory */
	mxFree(Xx);
	mxFree(XxT);
	mxFree(Xw);
	mxFree(XwT);
	mxFree(XxTXw);
	mxFree(Inv);
	mxFree(A);
	mxFree(AT);
	mxFree(XCoord);
	mxFree(YCoord);
	mxFree(W);
	mxFree(NoisyImgExt);
	mxFree(ResponsibilitiesExt);
	mxFree(Window);
	mxFree(WindowResp);
	mxFree(C11);
	mxFree(C12);
	mxFree(C22);
	mxFree(SqrtDetC);
	mxFree(C11Ext);
	mxFree(C12Ext);
	mxFree(C22Ext);
	mxFree(SqrtDetCExt);
	mxFree(WindowC11);
	mxFree(WindowC12);
	mxFree(WindowC22);
	mxFree(WindowSqrtDetC);
	mxFree(MyResp);
    
}    



/* Extend a 2D matrix (mirror extension at the edges) */
void Extend(double *Dest,double *Orig,int N, int M,int radius)
{
	int Offset,NdxCol,NdxRow,cnt;

	/* Copy the original matrix to the destination */
	Offset=radius*(N+2*radius)+radius;
	for(NdxCol=0;NdxCol<M;NdxCol++)
	{		
		memcpy(Dest+Offset+NdxCol*(N+2*radius),Orig+NdxCol*N,
			N*sizeof(double));
	}

	/* Extend to the left */
	for(NdxCol=0;NdxCol<radius;NdxCol++)
	{		
		memcpy(Dest+radius+NdxCol*(N+2*radius),Orig+(radius-NdxCol)*N,
			N*sizeof(double));
	}

	/* Extend to the right */
	Offset=(M+radius)*(N+2*radius)+radius;
	for(NdxCol=0;NdxCol<radius;NdxCol++)
	{		
		memcpy(Dest+Offset+NdxCol*(N+2*radius),Orig+(M-2-NdxCol)*N,
			N*sizeof(double));
	}

	/* Extend to the top */	
	for(NdxCol=0;NdxCol<M+2*radius;NdxCol++)
	{		
		for(NdxRow=0;NdxRow<radius;NdxRow++)
		{
			Dest[NdxRow+NdxCol*(N+2*radius)]=Dest[2*radius-NdxRow+NdxCol*(N+2*radius)];
		}		
	}

	/* Extend to the bottom */	
	for(NdxCol=0;NdxCol<M+2*radius;NdxCol++)
	{		
		for(NdxRow=radius+N,cnt=0;NdxRow<N+2*radius;NdxRow++,cnt++)
		{
			Dest[NdxRow+NdxCol*(N+2*radius)]=Dest[N+radius-2-cnt+NdxCol*(N+2*radius)];
		}		
	}

}
