/* mex-C: 
 * compute MAX1 maps for a *single* image at one scale
 *
 * Usage:
 * 
 *    [M1Map ARGMAX1 M1RowShift M1ColShift M1OriShifted] = 
 *      mexc_ComputeMAX1( S1Map, locationPerturbationRadius, 
 *      orientationPerturbationRadius, subsampleM1 );
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "mex.h"        /* the algorithm is connect to matlab */
#include <cmath>
#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 */
float **M1Map;                  /* MAX1 maps */
int **M1Trace;                  /* we currently do not allow shift in scale, so no need to track original scale */
const float** S1Map;            /* SUM1 maps */
int height, width;              /* size of MAX1 maps */
int nGaborOri;                  /* number of orientations for Gabor elements */
int numShift;                   /* number of local shifts, the same for all orientations and all scales */
int **rowShift, **colShift, **orientShift;  /* store the location/orientation shifts per scale and orientation */
int sizexSubsample, sizeySubsample;         /* (down sampled) size of M1 map */
int subsample;                              /* sub sampling step length */
int orientationPerturbationRadius;               /* allowed orientation perturbation */
int locationPerturbationRadius;             /* allowed location perturbation (number of pixels / subsampleM1) */

/* store all the shifts or perturbations in local maximum pooling */
void StoreShift()
{
    int orient, i, j, shift, ci, cj, si, sj, os;
    float alpha;

    /* store all the possible shifts for all orientations */
    numShift = (locationPerturbationRadius * 2 + 1) * (orientationPerturbationRadius*2+1);
    rowShift = (int**)mxCalloc( nGaborOri, sizeof(*rowShift) );
    colShift = (int**)mxCalloc( nGaborOri, sizeof(*colShift) );
    orientShift = (int**)mxCalloc( nGaborOri, sizeof(*orientShift) );

    for (orient=0; orient<nGaborOri; ++orient)  /* for each orientation */
    {
        rowShift[orient] = (int*)mxCalloc( numShift, sizeof(**rowShift) );
        colShift[orient] = (int*)mxCalloc( numShift, sizeof(**colShift) );
        orientShift[orient] = (int*)mxCalloc( numShift, sizeof(**orientShift) );
        
        /* order the shifts from small to large */
        alpha = (float)PI*orient/nGaborOri;
        ci = 0;
        for( i=0; i<=locationPerturbationRadius; ++i ) // location shifts
        {
            for (si=0; si<=1; si++) /* sign of i */
            {
                cj = 0;
                for( j = 0; j<=orientationPerturbationRadius; j++ ) // orientation shifts
                {
                    for (sj=0; sj<=1; sj++) /* sign of j */
                    {
                        shift = ci*(2*orientationPerturbationRadius+1) + cj; /* hybrid counter for location and orientation */
                        rowShift[orient][shift] = (int)floor(.5+i*(si*2-1)*subsample*cos(alpha));
                        colShift[orient][shift] = (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[orient][shift] = os;
                        if (j>0 || sj==1) /* triggers the orientation counter */
                        {
                            cj++;
                        }
                    }
                }
                if (i>0 || si==1)
                {
                    ci++;
                }
            }
        }
        
    }
}

void Compute()
{
    int orient, x, y;
    /* local maximum pooling at (x, y, orient) */
    float maxResponse, r;
    int shift, x1, y1, orient1, mshift;
    
    sizexSubsample = (int)floor((double)height / subsample);
    sizeySubsample = (int)floor((double)width / subsample);
    
    M1Map = (float**)mxCalloc( nGaborOri, sizeof(*M1Map) );
    M1Trace = (int**)mxCalloc( nGaborOri, sizeof(*M1Trace) );
    
    for (orient=0; orient<nGaborOri; orient++)
    {
        M1Map[orient] = (float*)mxCalloc( sizeySubsample * sizexSubsample, sizeof(**M1Map) );
        M1Trace[orient] = (int*)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; shift++ )
                {
                    x1 = x * subsample + rowShift[orient][shift];
                    y1 = y * subsample + colShift[orient][shift];
                    orient1 = orientShift[orient][shift];
                    if ((x1>=0)&&(x1<height)&&(y1>=0)&&(y1<width))
                    {
                        r = S1Map[orient1][x1+y1*height];
                        if (r>maxResponse)
                        {
                            maxResponse = r;
                            mshift = shift;
                        }
                    }
                }
                M1Map[orient][x+y*sizexSubsample] = maxResponse;
                M1Trace[orient][x+y*sizexSubsample] = (float)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;
    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: S1 maps
	 */
    pAS1Map = prhs[0];
    dims = mxGetDimensions(pAS1Map);
    nGaborOri = dims[0] * dims[1];
 
    S1Map = (const float**)mxCalloc( nGaborOri, sizeof(*S1Map) );   /* SUM1 maps */
    for (i=0; i<nGaborOri; ++i)
    {
        f = mxGetCell(pAS1Map, i);
        datatype = mxGetClassID(f);
        if (datatype != mxSINGLE_CLASS)
            mexErrMsgTxt("warning !! single precision required.");
        S1Map[i] = (const float*)mxGetPr(f);    /* get the pointer to cell content */
        height = mxGetM(f);    /* overwriting is ok, since it is constant */
        width = mxGetN(f);
    }
    
    /*
     * input variable 1: location shift radius
     */
    locationPerturbationRadius = (int)mxGetScalar(prhs[1]);
    
    /*
     * input variable 2: orientation shift radius
     */
    orientationPerturbationRadius = (int)mxGetScalar(prhs[2]);
    
    /*
     * input variable 3: sub sample step length
     */
    subsample = (int)mxGetScalar(prhs[3]);

    StoreShift();
    Compute();
    
    /* =============================================
     * Handle output variables.
     * ============================================= 
     */
    
    /*
     * output variable 0: M1 map
     */
    dimsOutput[0] = 1; dimsOutput[1] = nGaborOri;
	plhs[0] = mxCreateCellArray( 2, dimsOutput );
    dimsOutput[0] = sizexSubsample; dimsOutput[1] = sizeySubsample;
    for( i = 0; i < nGaborOri; ++i )
    {
        outPA = mxCreateNumericArray( 2, dimsOutput, mxSINGLE_CLASS, mxREAL );
        /* populate the real part of the created array */
        start_of_pr = (float*)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] = nGaborOri;
	plhs[1] = mxCreateCellArray( 2, dimsOutput );
    dimsOutput[0] = sizexSubsample; dimsOutput[1] = sizeySubsample;
    for( i = 0; i < nGaborOri; ++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] = nGaborOri; dimsOutput[1] = 1;
    plhs[2] = mxCreateCellArray( 2, dimsOutput );
    for( i = 0; i < nGaborOri; ++i )
    {
        dimsOutput[0] = numShift; dimsOutput[1] = 1;
        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] = nGaborOri; dimsOutput[1] = 1;
    plhs[3] = mxCreateCellArray( 2, dimsOutput );
    for( i = 0; i < nGaborOri; ++i )
    {
        dimsOutput[0] = numShift; dimsOutput[1] = 1;
        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] = nGaborOri; dimsOutput[1] = 1;
    plhs[4] = mxCreateCellArray( 2, dimsOutput );
    for( i = 0; i < nGaborOri; ++i )
    {
        dimsOutput[0] = numShift; dimsOutput[1] = 1;
        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 );
    }
}

