#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 SteeringMatrixSparse.c MatesLap.c Debugging.c lapack.a blas.a libf2c.a

Usage:
C = SteeringMatrixSparse(GradX, GradY, Resp, wsize, lambda, alpha);

Notes:
GradX=Gradient in the horizontal direction
GradY=Gradient in the vertical direction
Resp=Probability that a given pixel is not impulse corrupted
wsize=size of local analysis window
lambda=regularization for the elongation parameter
alpha=structure sensitive parameter
C=map of steering matrices

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,*GradX,*GradY,*C,*K,*MyResp;
	double *GradXExt,*GradYExt,*ResponsibilitiesExt,*G,*U;
	double S[2];
	double V[4];
	double lambda,alpha,aux,SumK,S1,S2;
	double AuxMat[4];
	double AuxMat2[4];
	int radius,wsize,M,N,NdxWin,NdxCol,NdxRow,NonZeros,NdxNonZero,NdxWinRow,NdxWinCol;
	int Dims[4];
	mxArray *My_plhs[1];
	mxArray *My_prhs[2];

    /* Get input data */      
	wsize=(int)mxGetScalar(prhs[3]);    
	lambda=mxGetScalar(prhs[4]);  
	alpha=mxGetScalar(prhs[5]);  
	GradX=mxGetPr(prhs[0]);
	GradY=mxGetPr(prhs[1]);
	Responsibilities=mxGetPr(prhs[2]);
    
    
    /* Create output array */
	N=mxGetM(prhs[0]);  /* Number of rows in the input image */
	M=mxGetN(prhs[0]);  /* Number of columns in the input image */
	Dims[0]=2;
	Dims[1]=2;
	Dims[2]=N;
	Dims[3]=M;
    plhs[0]=mxCreateNumericArray(4,Dims,mxDOUBLE_CLASS,mxREAL);
	C=mxGetPr(plhs[0]);
	    
    
    /* Get working variables */
	if (wsize%2==0)
	{
		wsize++;
	}
	radius = wsize/2;

	/* Get the filter */
	My_prhs[0]=mxCreateString("disk");
	My_prhs[1]=mxCreateDoubleMatrix(1,1,mxREAL);
	(*mxGetPr(My_prhs[1]))=(double)radius;
	mexCallMATLAB(1,My_plhs,2,My_prhs,"fspecial");
	K=mxGetPr(My_plhs[0]);
	aux=K[radius*wsize+radius];
	SumK=0.0;
	for(NdxWin=0;NdxWin<wsize*wsize;NdxWin++)
	{
		K[NdxWin]/=aux;
		SumK+=K[NdxWin];
	}
	

    /* Prepare auxiliary arrays */
	GradXExt=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));
	GradYExt=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));
	ResponsibilitiesExt=mxMalloc((N+2*radius)*(M+2*radius)*sizeof(double));
	G=mxMalloc((wsize*wsize)*2*sizeof(double));
	U=mxMalloc((wsize*wsize)*2*sizeof(double));
	MyResp=mxMalloc(wsize*wsize*sizeof(double));

	/* Extend the gradient maps and the responsibility map */
	Extend(GradXExt,GradX,N,M,radius);	
	Extend(GradYExt,GradY,N,M,radius);	
	Extend(ResponsibilitiesExt,Responsibilities,N,M,radius);
	

	/* Estimate the steering matrices */
	for(NdxCol=0;NdxCol<M;NdxCol++)
	{
		for(NdxRow=0;NdxRow<N;NdxRow++)
		{

			/* Get the responsibilities in MyResp */
			for(NdxWin=0;NdxWin<wsize;NdxWin++)
			{
				memcpy(MyResp+NdxWin*wsize,ResponsibilitiesExt+(NdxCol+NdxWin)*(N+2*radius)+NdxRow,
					wsize*sizeof(double));
			}

			/* Count the number of pixels with nonzero responsibility */
			NonZeros=0; 
			for(NdxWin=0;NdxWin<wsize*wsize;NdxWin++)
			{
				if (MyResp[NdxWin]!=0.0)
				{
					NonZeros++;
				}
			}

			/* Get the window with the relevant gradient values */
			NdxNonZero=0;
			for(NdxWin=0;NdxWin<wsize*wsize;NdxWin++)
			{
				/* Process only pixels with nonzero responsibility */
				if (MyResp[NdxWin]!=0.0)				
				{
					/* Coordinates inside the window */
					NdxWinCol=NdxWin/wsize;
					NdxWinRow=NdxWin%wsize;

					/* Get the gradient value and multiply it by the filter */
					G[NdxNonZero]=K[NdxWin]*GradXExt[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];
					G[NdxNonZero+NonZeros]=K[NdxWin]*GradYExt[(NdxCol+NdxWinCol)*(N+2*radius)+NdxRow+NdxWinRow];
					NdxNonZero++;
				}
			}


			/* Compute Singular Value Decomposition */
			EconomySVD(G,S,U,V,NonZeros,2);

			
			/* Regularized values */
			S1 = (S[0] + lambda) / (S[1] + lambda);
			S2 = (S[1] + lambda) / (S[0] + lambda);
			
			/* Compute the steering matrix */			
			MatrixProduct(V,V,AuxMat,2,1,2);
			ScalarMatrixProduct(S1,AuxMat,AuxMat,2,2);
			MatrixProduct(V+2,V+2,AuxMat2,2,1,2);
			ScalarMatrixProduct(S2,AuxMat2,AuxMat2,2,2);
			MatrixSum(AuxMat,AuxMat2,AuxMat,2,2);	

			/* Apply adaptive weighting and write on output */
			aux=pow((S[0] * S[1] + 0.0000001) / SumK, alpha);
			ScalarMatrixProduct(aux*Responsibilities[NdxCol*N+NdxRow],AuxMat,
				C+4*(NdxCol*N+NdxRow),2,2);

		}

	}

	/* Release memory */
	mxDestroyArray(My_plhs[0]);
	mxDestroyArray(My_prhs[0]);
	mxDestroyArray(My_prhs[1]);
	mxFree(GradXExt);
	mxFree(GradYExt);
	mxFree(ResponsibilitiesExt);
	mxFree(G);
	mxFree(U);
	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)];
		}		
	}

}
