/* mex-C: 
 * compute MAX1 maps for a single image
 *
 * Usage:
 * 
 *    [M1Map M1Trace M1RowShift M1ColShift M1OriShift] = 
 *      CcomputeMAX1( nGaborOri, S1Map, locationPerturbationFraction, 
 *      orientationPerturbationRadius, sizeGaborPerScale, subsampleM1 );
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include "mex.h"        /* the algorithm is connect to matlab */
#include "math.h"
#include "matrix.h"
#define ABS(x) ((x)>0? (x):(-(x)))
#define MAX(x, y) ((x)>(y)? (x):(y))
#define MIN(x, y) ((x)<(y)? (x):(y))
# define PI 3.1415926
# define NEGMAX -1e10

/* variable declaration */
double **M1Map;                  /* MAX 1 maps */
int **M1Trace;                  /* we currently do not allow shift in scale, so no need to track original scale */
const double** S1Map;            /* SUM1 maps */
int height, width;              /* size of MAX1 maps */
int nGaborScale;                /* number of Gabor scales */
int nGaborOri;                  /* number of orientations for Gabor elements */
double *sizeGabor;               /* size of Gabor filters at different scales */ 
int *numShift;                  /* number of local shifts per scale */
int **rowShift, **colShift, **orientShift;  /* store the location/orientation shifts per scale */
int sizexSubsample, sizeySubsample; /* (down sampled) size of M1 map */
int *numLocationShift;              /* number of location shifts per scale */
int subsample;                      /* sub sampling step length */
int orientationPerturbRadius;       /* allowed orientation perturbation */
double locationPerturbRadius;        /* allowed location perturbation, relative to the Gabor filter size */

/* store all the shifts or perturbations in local maximum pooling */
void StoreShift()  
{
    int orient, i, j, shift, ci, cj, si, sj, os, iS;
    double alpha;
    int locationShiftLimit;
    
    /* store all the possible shifts for all orientations */
    
    numShift = mxCalloc( nGaborScale, sizeof(*numShift) );
    rowShift = mxCalloc( nGaborScale, sizeof(*rowShift) );
    colShift = mxCalloc( nGaborScale, sizeof(*colShift) );
    orientShift = mxCalloc( nGaborScale, sizeof(*colShift) );
    numLocationShift = mxCalloc( nGaborScale, sizeof(*numLocationShift) );
    
    for( iS = 0; iS < nGaborScale; ++iS ) /* for each scale */
    {
        locationShiftLimit = (int)floor( sizeGabor[iS] * locationPerturbRadius / subsample );
       
        numLocationShift[iS] = (locationShiftLimit*2+1);
        numShift[iS] = numLocationShift[iS]*(orientationPerturbRadius*2+1);
        rowShift[iS] = mxCalloc( nGaborOri * numShift[iS], sizeof(**rowShift) );
        colShift[iS] = mxCalloc( nGaborOri * numShift[iS], sizeof(**colShift) );
        orientShift[iS] = mxCalloc( nGaborOri * numShift[iS], sizeof(**orientShift) );
        
        /* order the shifts from small to large */
        for (orient=0; orient<nGaborOri; orient++)
        {
            alpha = (double)PI*orient/nGaborOri;
            ci = 0;
            for( i=0; i<=locationShiftLimit; ++i )
            {
                for (si=0; si<=1; si++) /* sign of i */
                {
                    cj = 0;
                    for( j = 0; j<=orientationPerturbRadius; j++ )
                    {
                        for (sj=0; sj<=1; sj++) /* sign of j */
                        {
                            shift = ci*(2*orientationPerturbRadius+1) + cj; /* hybrid counter for location and orientation */
                            rowShift[iS][shift+orient*numLocationShift[iS]] = (int)floor(.5+i*(si*2-1)*subsample*cos(alpha));
                            colShift[iS][shift+orient*numLocationShift[iS]] = (int)floor(.5+i*(si*2-1)*subsample*sin(alpha));
                            os = orient + j*(sj*2-1);
                            if (os<0)
                            {
                                os += nGaborOri;
                            }
                            else if (os>=nGaborOri)
                            {
                                os -= nGaborOri;
                            }
                            orientShift[iS][shift+orient*numLocationShift[iS]] = os;
                            if (j>0 || sj==1) /* triggers the orientation counter */
                            {
                                cj++;
                            }
                        }
                    }
                    if (i>0 || si==1)
                    {
                        ci++;
                    }
                }
            }
        }
        
    }
}

void Compute()
{
    int orient, x, y, iS;
    /* local maximum pooling at (x, y, orient) */
    double maxResponse, r;
    int shift, x1, y1, orient1, mshift;
    
    sizexSubsample = (int)floor(height / subsample);
    sizeySubsample = (int)floor(width / subsample);
    
    M1Map = mxCalloc( nGaborOri * nGaborScale, sizeof(*M1Map) );
    M1Trace = mxCalloc( nGaborOri * nGaborScale, sizeof(*M1Trace) );
    
    for( iS = 0; iS < nGaborScale; ++iS )
    {
        for (orient=0; orient<nGaborOri; orient++)
        {
            M1Map[iS*nGaborOri+orient] = mxCalloc( sizeySubsample * sizexSubsample, sizeof(**M1Map) );
            M1Trace[iS*nGaborOri+orient] = mxCalloc( sizeySubsample * sizexSubsample, sizeof(**M1Trace) );
            
            for (y=0; y<sizeySubsample; y++)
                for (x=0; x<sizexSubsample; x++)
                {
                    maxResponse = NEGMAX; mshift = 0;
                    for( shift=0; shift<numShift[iS]; shift++ )
                    {
                        x1 = x * subsample + rowShift[iS][shift+orient*numLocationShift[iS]];
                        y1 = y * subsample + colShift[iS][shift+orient*numLocationShift[iS]];
                        orient1 = orientShift[iS][shift+orient*numLocationShift[iS]];
                        if ((x1>=0)&&(x1<height)&&(y1>=0)&&(y1<width))
                        {
                            r = S1Map[orient1][x1+y1*height];
                            if (r>maxResponse)
                            {
                                maxResponse = r;
                                mshift = shift;
                            }
                        }
                    }
                    M1Map[iS*nGaborOri+orient][x+y*sizexSubsample] = maxResponse;
                    M1Trace[iS*nGaborOri+orient][x+y*sizexSubsample] = (double)mshift;
                }
        }
    }
    
}


/* entry point */
void mexFunction(int nlhs, mxArray *plhs[], 
                 int nrhs, const mxArray *prhs[])                
{
    int ind, i, x, y, o, dataDim, j, bytes_to_copy, nGaborFilter;
    const mxArray *f;
    const mxArray *pAS1Map;
    mxArray *outPA;
    mwSize ndim;
    const mwSize* dims;
    mwSize dimsOutput[2];
    void* start_of_pr;
    mxClassID datatype;

	/*
	 * input variable 0: nGaborOri
	 */
	nGaborOri = (int)mxGetScalar(prhs[0]);
 
    /*
	 * input variable 1: S1 maps
	 */
    pAS1Map = prhs[1];
    dims = mxGetDimensions(pAS1Map);
    nGaborFilter = dims[0] * dims[1];
    nGaborScale = nGaborFilter / nGaborOri;
 
    S1Map = mxCalloc( nGaborFilter, sizeof(*S1Map) );   /* SUM1 maps */
    for (i=0; i<nGaborFilter; ++i)
    {
        f = mxGetCell(pAS1Map, i);
        datatype = mxGetClassID(f);
        if (datatype != mxDOUBLE_CLASS)
            mexErrMsgTxt("warning !! double precision required.");
        S1Map[i] = (const double*)mxGetPr(f);    /* get the pointer to cell content */
        height = mxGetM(f);    /* overwriting is ok, since it is constant */
        width = mxGetN(f);
    }
    
    /*
     * input variable 2: location shift radius
     */
    locationPerturbRadius = (double)mxGetScalar(prhs[2]);
    
    /*
     * input variable 3: orientation shift radius
     */
    orientationPerturbRadius = (int)mxGetScalar(prhs[3]);
    
    /*
     * input variable 4: sizes of Gabor filters
     */
    sizeGabor = (double*)mxGetPr(prhs[4]);
    
    /*
     * input variable 5: sub sample step length
     */
    subsample = (int)mxGetScalar(prhs[5]);

    StoreShift();
    Compute();
    
    /* =============================================
     * Handle output variables.
     * ============================================= 
     */
    
    /*
     * output variable 0: M1 map
     */
    dimsOutput[0] = 1; dimsOutput[1] = nGaborScale * nGaborOri;
	plhs[0] = mxCreateCellArray( 2, dimsOutput );
    dimsOutput[0] = sizexSubsample; dimsOutput[1] = sizeySubsample;
    for( i = 0; i < nGaborOri * nGaborScale; ++i )
    {
        outPA = mxCreateNumericArray( 2, dimsOutput, mxDOUBLE_CLASS, mxREAL );
        /* populate the real part of the created array */
        start_of_pr = (double*)mxGetData(outPA);
        bytes_to_copy = dimsOutput[0] * dimsOutput[1] * mxGetElementSize(outPA);
        memcpy( start_of_pr, M1Map[i], bytes_to_copy );
        mxSetCell( plhs[0], i, outPA );
    }
    
    /*
     * output variable 1: M1 trace
     */
    dimsOutput[0] = 1; dimsOutput[1] = nGaborScale * nGaborOri;
	plhs[1] = mxCreateCellArray( 2, dimsOutput );
    dimsOutput[0] = sizexSubsample; dimsOutput[1] = sizeySubsample;
    for( i = 0; i < nGaborOri * nGaborScale; ++i )
    {
        outPA = mxCreateNumericArray( 2, dimsOutput, mxINT32_CLASS, mxREAL );
        /* populate the real part of the created array */
        start_of_pr = (int*)mxGetData(outPA);
        bytes_to_copy = dimsOutput[0] * dimsOutput[1] * mxGetElementSize(outPA);
        memcpy( start_of_pr, M1Trace[i], bytes_to_copy );
        mxSetCell( plhs[1], i, outPA );
    }
    
    /*
     * output variable 2: stored Gabor shifts : row shifts
     */
    dimsOutput[0] = 1; dimsOutput[1] = nGaborScale;
    plhs[2] = mxCreateCellArray( 2, dimsOutput );
    for( i = 0; i < nGaborScale; ++i )
    {
        dimsOutput[0] = 1; dimsOutput[1] = numShift[i];
        outPA = mxCreateNumericArray( 2, dimsOutput, mxINT32_CLASS, mxREAL );
        start_of_pr = (int*)mxGetData(outPA);
        bytes_to_copy = dimsOutput[0] * dimsOutput[1] * mxGetElementSize(outPA);
        memcpy( start_of_pr, rowShift[i], bytes_to_copy );
        mxSetCell( plhs[2], i, outPA );
    }
    
    /*
     * output variable 3: stored Gabor shifts : col shifts
     */
    dimsOutput[0] = 1; dimsOutput[1] = nGaborScale;
    plhs[3] = mxCreateCellArray( 2, dimsOutput );
    for( i = 0; i < nGaborScale; ++i )
    {
        dimsOutput[0] = 1; dimsOutput[1] = numShift[i];
        outPA = mxCreateNumericArray( 2, dimsOutput, mxINT32_CLASS, mxREAL );
        start_of_pr = (int*)mxGetData(outPA);
        bytes_to_copy = dimsOutput[0] * dimsOutput[1] * mxGetElementSize(outPA);
        memcpy( start_of_pr, colShift[i], bytes_to_copy );
        mxSetCell( plhs[3], i, outPA );
    }
    
    /*
     * output variable 4: stored Gabor shifts : orientation shifts
     */
    dimsOutput[0] = 1; dimsOutput[1] = nGaborScale;
    plhs[4] = mxCreateCellArray( 2, dimsOutput );
    for( i = 0; i < nGaborScale; ++i )
    {
        dimsOutput[0] = 1; dimsOutput[1] = numShift[i];
        outPA = mxCreateNumericArray( 2, dimsOutput, mxINT32_CLASS, mxREAL );
        start_of_pr = (int*)mxGetData(outPA);
        bytes_to_copy = dimsOutput[0] * dimsOutput[1] * mxGetElementSize(outPA);
        memcpy( start_of_pr, orientShift[i], bytes_to_copy );
        mxSetCell( plhs[4], i, outPA );
    }
}

