#include "mex.h"
#include <math.h>
#include <matrix.h>

/* actual nms function */
void nms(double *x, double *y, double *mask, double *img, int rows, int cols, int nms_region, double thresh, double minpix, int mask_rgn)
{
   int i, j, m, n, k, l;  /* iteration variables */
   int cur_point, xs, xe, ys, ye;   /* indices for the matrix */
   int islocalmax(double *, int, int, int, int, int, int, int, int, double, double);
   int *mark = mxMalloc(rows * cols * sizeof(int));
   double dist1, dist2;
   
   /*mexPrintf("The search region is %d.\n", 2*nms_region+1);*/
   
   cur_point = 1;
   for(i=0; i<rows*cols; i++)
	mark[i] = 0;

   /* for testing purpose only */
   /*rows = 55;	cols = 55;*/

   /* for each pixel in the image, see if it is local maximum */
   for(j=1; j<=cols; j++)
   {
	for(i=1; i<=rows; i++)
	{
		if(mark[i-1+rows*(j-1)] == 2)	/* already marked as non-maximal */
			continue;

		xs = i - nms_region;
		if(xs < 1)
			xs = 1;

		ys = j - nms_region;
		if(ys < 1)
			ys = 1;

		xe = i + nms_region;
		if(xe > rows)
			xe = rows;

		ye = j + nms_region;
		if(ye > cols)
			ye = cols;

		if(islocalmax(img, i, j, xs, ys, xe, ye, rows, cols, thresh, minpix) == 1) {
			x[cur_point] = i;
			y[cur_point] = j;
			cur_point ++;

			for(m=xs; m<=xe; m++)
			for(n=ys; n<=ye; n++)
				mark[m-1+rows*(n-1)]=2;

			mark[i-1+rows*(j-1)]=1;		
		}
		else
			mark[i-1+rows*(j-1)] = 2;	
	}
   }

  /* now we have a list of particle coordinates, we are ready to create the mask matrix */
  for(i=1; i<cur_point; i++)
  {
	m = x[i];	n = y[i];		/* coordinate of the (i+1)th particle */
	xs = m - mask_rgn;
	if(xs < 1)	xs = 1;

	ys = n - mask_rgn;
	if(ys < 1)	ys = 1;

	xe = m + mask_rgn;
	if(xe > rows)	xe = rows;
	
	ye = n + mask_rgn;
	if(ye > cols) 	ye = cols;

	for(j=xs; j<=xe; j++)
	for(k=ys; k<=ye; k++)
	{
	   if(mask[j-1+rows*(k-1)] == 0)	/* pixel not claimed yet */
		mask[j-1+rows*(k-1)] = i;
	   else {				/* pixel already claimed */
		l = mask[j-1+rows*(k-1)];

		dist1 = sqrt(pow(j-m, 2) + pow(k-n, 2));
		dist2 = sqrt(pow(j-x[l], 2) + pow(k-y[l], 2));

		if(dist1 < 1.4 || dist1 <= (dist2 - 2))
			mask[j-1+rows*(k-1)] = i;
		else if((dist1 - dist2 < 2) && (dist1 - dist2 > -2))
			mask[j-1+rows*(k-1)] = 0;
	 	}
	}			
  }  

	/* free memory */
   mxFree(mark);

   return;
}

/* subroutine that determines if point i, j is the local maxima in a search region of nms_region radius */
int islocalmax(double *img, int x, int y, int xs, int ys, int xe, int ye, int rows, int cols, double thresh, double minpix)
{
   int i, j, temp, divider=0, positive=0;	
   double around[20];
   double local_ave, rms, a1_ave, a2_ave, cur_pval; 
   double corner_ave(double *, int, int, int, int, int);

   /*mexPrintf("Testing point (%d, %d) = %.0f; search area = (%d, %d) - (%d, %d)...\n", x, y, img[x-1+(y-1)*rows], xs, ys, xe, ye);*/
   temp = img[x-1+(y-1)*rows];

   /* any of the pixels greater than center pixel return 0 */
   for(i=xs; i<=xe; i++)
   for(j=ys; j<=ye; j++)
   {
	if(img[i-1+(j-1)*rows] > temp)
		return 0;
	else {
	    if(i<=x+1 && i>=x-1 && j<=y+1 && j>=y-1)  		/* surrounding pixels */
		around[++divider] = img[i-1+(j-1)*rows];
	    }
   }
   
   /* now center pixel is the brightest, make sure it is not a spike pixel.
      this is achieved by making sure it has sufficient surrounding pixels that stands above the local average */
   a1_ave = corner_ave(img, x, y, rows, cols, 1);
   a2_ave = corner_ave(img, x, y, rows, cols, 2);
   local_ave = (corner_ave(img, x, y, rows, cols, 3) + corner_ave(img, x, y, rows, cols, 4))/2.0; 
   
   if(temp < local_ave + thresh)
	return 0;

   /*mexPrintf("Local average is %5.1f\n", local_ave);*/
   for(i=0; i<divider; i++)
	if(around[i] > local_ave)	
		positive ++;
 
   if(positive >= minpix && a1_ave > local_ave) {
	return 1;
   }
   else
	return 0;
}


/* sub routine that calculates corner pixel averages */
double corner_ave(double *img, int x, int y, int rows, int cols, int offset)
{
   int x_ne, x_se, x_nw, x_sw;
   int y_ne, y_se, y_nw, y_sw;
   int divider = 4;
   double sum = 0;

   x_ne = x + offset;
   y_ne = y + offset;  
   if(x_ne > rows || y_ne > cols)
	divider--;
   else
	sum += img[x_ne-1+rows*(y_ne-1)];

   x_se = x + offset;
   y_se = y - offset;
   if(x_se > rows || y_se < 1)
	divider --;
   else
	sum += img[x_se-1+rows*(y_se-1)];

   x_nw = x - offset;
   y_nw = y + offset;
   if(x_nw < 1 || y_nw > cols)
	divider --;
   else
	sum += img[x_nw-1+rows*(y_nw-1)];

   x_sw = x - offset;
   y_sw = y - offset;
   if(x_sw < 1 || y_sw < 1)
	divider --;
   else
	sum += img[x_sw-1+rows*(y_sw-1)];

   return sum/divider;
}


/* entry point mex function */
void mexFunction( int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[] )
{
   double *x, *y, *img, *mask, thresh, minpix;
   int nms_region, initial_length, mask_rgn;
   mwSize cols, rows;
   void nms(double *, double *, double *, double *, int, int, int, double, double, int);
   
   /* mexPrintf("Number of outputs: %d", nlhs); */
   
   /* check the correct number of inputs and outputs */
   if (nrhs != 5) {
       mexErrMsgTxt("Incorrect number of inputs (required: image, nms_region, thresh, minpix).\n");  }
   else if (nlhs != 3) {           /* do not use nlhs != 1 here since in command line nlhs = 0 */
       mexErrMsgTxt("Incorrect number of outputs. Should be [x, y, mask].\n"); 
   }
   
   /* Assign pointers input trajectory array */
   img = mxGetPr(prhs[0]);
   rows = mxGetM(prhs[0]);      /* number of particles */
   cols = mxGetN(prhs[0]);      
   nms_region = *mxGetPr(prhs[1]);
   thresh = *mxGetPr(prhs[2]);	/* threshold (e.g. 3*rms) */
   minpix = *mxGetPr(prhs[3]);   /* number of pixels in center above threshold */
   mask_rgn = *mxGetPr(prhs[4]);
   /* note: although img is MxN, it now appears as a 1x(MxN) or linear vector */

   /* check the input */
   /*mexPrintf("\nThe first pixel in input image is %f.\n", img[0]);*/
	
   /* Allocate memory for output array */
   initial_length = (int) (rows * cols / 10);
   plhs[0] = mxCreateDoubleMatrix(initial_length, 1, mxREAL);
   plhs[1] = mxCreateDoubleMatrix(initial_length, 1, mxREAL);
   plhs[2] = mxCreateDoubleMatrix(rows, cols, mxREAL);
    
   x = mxGetPr(plhs[0]);
   y = mxGetPr(plhs[1]);
   mask = mxGetPr(plhs[2]);
   
   nms(x, y, mask, img, rows, cols, nms_region, thresh, minpix, mask_rgn);
}
