/*
 * Usage:
 *    [DT1 DT2 DT3] = mexc_find_S3(paths,featureFolder,dictionaries,neighborhoods);
 *
 *	"dictionaries" is a 3 by 1 cell array. Each cell is a cell array itself, which contains
		templates. dictionaries{1} is empty (atomic). dictionaries{2} are 
		S2 templates (active basis). dictionary{3} are S3 templates.
		
	A most basic element is a 3 tuple (x,y,channel). A SparseTemplate is a list of 3 tuples
	and represented in matlab mxArray is a 3 by #element matrix.
	
 *	"paths" is a 2D cell array of strings, which contains folder hierarchy and image name, 
 * 		so that its path for S1M1/S2M3/S3 maps can be found. The path is relative to the folder
		where this file resides. Each path is a cell array of strings. The first item is
		the root directory for image; the last item is the image name.
 *	"featurFolder" is the folder where feature maps are stored. Feature maps are organized 
		the same as the source images.
 */
 
/*
			TODO
- Use smart pointers e.g. auto_ptr.
*/

#include "HierTemplate.h"
#include "Log.h"
#include <string>
#include <stdio.h>
#include <stdlib.h>
#include "mex.h"
#include "mat.h"
#include "math.h"
#include <ctime>
#define PI 3.1415926
#define ABS(x) ((x)>0? (x):(-(x)))
#define MAX(x, y) ((x)>(y)? (x):(y))
#define MIN(x, y) ((x)<(y)? (x):(y))
#define ROUND(x) (floor((x)+.5))

LOG_DECLARE; // A global variable named __log is created

typedef float Pixel_t;		/* change "float" to double if double precision is used */
typedef float Param_t;
typedef vector<string> Path_t;	/* a path is a vector of strings indicating the folder hierarchy
								with the last item being the file name */
typedef SparseTemplate<Param_t> SparseT_t;
typedef HierTemplate<Pixel_t,Param_t> HierT_t;
typedef HierT_t::Dictionaries_t Dict_t;
								
/* global variables */
vector<Path_t> paths;
Path_t featureFolder;
const bool displayDebugInfo = false;

/* The following maps only store feature maps for a single image. */
vector<const Pixel_t**> sumMapAllLevel, maxMapAllLevel;
vector<const int**> argmaxMapAllLevel;
/* The following maps or lists are shared among all image. */
// vector<const int**> templateAdjacencyAllLevel; // this is used in computing MAX maps, not yet used here (to improve)
/**/
int currentWidth, currentHeight, nChS3;

/* About indexing of levels: 
	dicts[0] is deliberately left empty, since there is no S0 template.
		S1 template is a simple active basis model.
	xxx_map[0] means the source image. sum_map_[1] means SUM1 map.
	
	The matlab data format of dictionaries is a 2D cell array of 3-tuples.
*/


/* convert an matlab mxArray (3 by #elements) to SparseTemplate */
void convertMxArrayToSparseTemplate( const mxArray* ma, SparseT_t* sp )
{
	assert( mxGetM(ma) == 3 ); // The mxArray should have exactly 3 rows.
	const Param_t* f = (const Param_t*)mxGetPr( ma );

    if( f == NULL )
    {
        mexErrMsgTxt("in convertMxArrayToSparseTemplate: const Param_t* f==NULL");
    }

	int nF = mxGetN(ma);
	vector< vector<Param_t> > elements;
	for( int iF = 0; iF < nF; ++iF )
	{
		vector<Param_t> vv;
		vv.push_back( f[0+iF*3] );
		vv.push_back( f[1+iF*3] );
		vv.push_back( f[2+iF*3] );
		elements.push_back( vv );
	}

	sp->set_elements(elements);
}

/* The input should be a cell array of a cell array.
	So that in matlab, input{2}{10} means the 10th SparseTemplate
	in the 2nd level dictionary. And input{2}{10} would be 
	a 3 by n matrix of scalars.
 */
void inputDictionaries( const mxArray* cpA, Dict_t &dicts )
{
	mxClassID datatype;
	if( mxGetNumberOfElements(cpA) != 3 )
	{
		mexErrMsgTxt("warning !! dictionaries should be of length 3");
	}
    
	/* insert an empty layer of dictionary */
    dicts.clear();
    vector<const SparseT_t*> tmp;
    dicts.push_back( tmp );
 
    for( int i = 0; i <= 2; ++i )
    {
		const mxArray* cpA2 = mxGetCell( cpA, i );
		int nT = mxGetNumberOfElements( cpA2 );
		if( displayDebugInfo )
		{
			mexPrintf( "S%d dictionary has %d templates.\n", i+1, nT );
			mexEvalString("drawnow");
		}
		
		vector<const SparseT_t*> v( 0 );
		for( int j = 0; j < nT; ++j )
		{
			const mxArray* cpA3 = mxGetCell( cpA2, j );
			if( cpA3 == NULL )
			{
				v.push_back( NULL );
				continue;
			}
		
			SparseT_t* pS = new SparseT_t();
		
			
			datatype = mxGetClassID( cpA3 );
			if( datatype != mxSINGLE_CLASS )
			{
				mexErrMsgTxt("warning !! single precision required.");
			}
			convertMxArrayToSparseTemplate( cpA3, pS );
			if( displayDebugInfo )
			{
            	mexPrintf( "level %d entry %d, %d elements\n", i, j, pS->elements().size() );
            }
			v.push_back( pS );  // remember to delete pS
		}
		dicts.push_back(v);
	}
}

/* This function accompanies inputDictionaries */
void destroyDictionaries( Dict_t &dicts )
{
	for( Dict_t::iterator it = dicts.begin(); it != dicts.end(); ++it )
	{
		for( vector<const SparseT_t*>::iterator jt = it->begin(); jt != it->end(); ++jt )
		{
            if( *jt != NULL )
            {
                delete (*jt);
            }
		}
	}
}

template <typename T>
const T** convert2DConstMxArrayToCArray( const mxArray* ma, /* output: */int* width=NULL, int* height=NULL, int* nCh=NULL )
{
	/* ma is the pointer to a 2D mxArray */
	int nch = mxGetNumberOfElements( ma );
	const T** pp = (const T**)mxCalloc( nch, sizeof(*pp) );
	if( width!=NULL && height!=NULL && nCh!=NULL )
	{
		*width = mxGetN( mxGetCell(ma,0) );
		*height = mxGetM( mxGetCell(ma,0) );
		*nCh = nch;
	}
	for( int c = 0; c < nch; ++c )
	{
		const T* p = (const T*) mxGetPr( mxGetCell(ma,c) );
        if( p == NULL )
        {
            mexErrMsgTxt( "converting 2d const mxarray, p == NULL" );
        }
		pp[c] = p;
	}
	return pp;
}

template <typename T>
const T** load2DConstMxArrayToCArray( MATFile* mf, string varName, int* width=NULL, int* height=NULL, int* nCh=NULL )
{
	const mxArray* ma = matGetVariable(mf, varName.c_str());
    if( ma == NULL )
    {
        mexPrintf( "%s\n", varName.c_str() );
        mexErrMsgTxt( "variable not found" );
    }
	const T** a = convert2DConstMxArrayToCArray<T>( ma, width, height, nCh );
	return a;
}



/* compute, save and load feature maps */
void prepareFeatureMap( const string& featurePath, const Dict_t& dicts ) // level 1 means S1/M1/ARGMAX1 map
{
	MATFile *mf;
	
	/* initialize the feature maps :
		A dumb element is inserted as the first element, so that xxxMapXXX[1] is the first level.
		There is no need to delete the pointers in the following containers.
	*/
	const Pixel_t** tmp = NULL;
	sumMapAllLevel.clear();
	sumMapAllLevel.push_back( tmp );
	maxMapAllLevel.clear();
	maxMapAllLevel.push_back( tmp );
	argmaxMapAllLevel.clear();
	const int** tmp2 = NULL;
	argmaxMapAllLevel.push_back( tmp2 );

	mf = matOpen(featurePath.c_str(), "u" );
	if( mf == NULL )
    {
        mexPrintf("%s\n", featurePath.c_str());
        mexErrMsgTxt( "error openning file" );
    }
    
	const mxArray* ma;
	int nch;
	const Pixel_t** pp;
	const int** ppi;
	const Param_t* ppa;
	
	// deal with level 1
	/* MAX1 map */
	pp = load2DConstMxArrayToCArray<Pixel_t>( mf, "MAX1map" );
    if( pp == NULL )
    {
        mexErrMsgTxt( "error loading max1 map" );
    }
	maxMapAllLevel.push_back( pp );
    
    
	/* ARGMAX1 map */
	ppi = load2DConstMxArrayToCArray<int>( mf, "ARGMAX1map" );
    if( ppi == NULL )
    {
        mexErrMsgTxt( "error loading argmax1 map" );
    }
	argmaxMapAllLevel.push_back( ppi );
	/* drow,dcol,dchn */
	//ppa = (const Param_t*)mxGetPr(  );
	
	// deal with level 2
	// prepare sum2, max2, argmax2 map
	/* MAX1 map */
	pp = load2DConstMxArrayToCArray<Pixel_t>( mf, "MAX2map" );
    if( pp == NULL )
    {
        mexErrMsgTxt( "error loading max2 map" );
    }
	maxMapAllLevel.push_back( pp );
	/* ARGMAX1 map */
	ppi = load2DConstMxArrayToCArray<int>( mf, "ARGMAX2map" );
    if( ppi == NULL )
    {
        mexErrMsgTxt( "error loading argmax2 map" );
    }
	argmaxMapAllLevel.push_back( ppi );
    
	/* drow,dcol,dchn */

	// deal with level 3
	/* SUM3 map */
	pp = load2DConstMxArrayToCArray<Pixel_t>( mf, "SUM3map",/* following is the output : */ &currentWidth, &currentHeight, &nChS3 );
    if( pp == NULL )
    {
        mexErrMsgTxt( "error loading sum3 map" );
    }
	sumMapAllLevel.resize(3); /* fill in previous levels with dumb objects */
	sumMapAllLevel.push_back( pp );
    argmaxMapAllLevel.push_back(NULL);
    maxMapAllLevel.push_back(NULL);
    
	matClose(mf);
}


void compute( const Dict_t &dicts, const Dict_t &nb,
		vector<SparseT_t>& deformedTemplatesLevel1,
		vector<SparseT_t>& deformedTemplatesLevel2,
		vector<SparseT_t>& deformedTemplatesLevel3 )
{
	/* check and load feature maps */
	// check if S1 M1 maps are stored
	// if sum2 max2 maps are not stored
	// 		load sum2 templates
	// 		compute sum2 maps
	//		compute max2 and argmax2
	// load S3 templates
	// compute sum3 map

	deformedTemplatesLevel1.clear();
	deformedTemplatesLevel2.clear();
	deformedTemplatesLevel3.clear();
	
	int nPath = paths.size();
	string feaRoot = featureFolder[0]; // featureFolder may contain hierarchy
	for( int j = 1; j < featureFolder.size(); ++j )
	{
		feaRoot += "/" + featureFolder[j];
	}
    
    
	for( int i = 0; i < nPath; ++i )
	{
		string s = feaRoot; // featureFolder may contain hierarchy
		for( int j = 1; j < paths[i].size()-1; ++j ) // image relative path starts from the 2nd item
													// the last item is the image name
		{
			s += "/" + paths[i][j];
		}
		s += "/" + paths[i][paths[i].size()-1];
		
		// timer:
		clock_t init, final;
		init=clock();
		
		prepareFeatureMap( s, dicts );
        final=clock()-init;
        if( displayDebugInfo )
        {
        	mexPrintf( "(image %d) prepared feature map: %.3f sec.\n", i, (double)final / ((double)CLOCKS_PER_SEC) );
			mexPrintf( "prepared feature map: %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
			mexEvalString("drawnow");
        }
		
		/* perform inference */
		// re-point to new feature maps
		HierT_t ht(0,0,0,3);
		ht.set_all_sum_map( sumMapAllLevel );
		ht.set_all_max_map( maxMapAllLevel );
		ht.set_all_argmax_map( argmaxMapAllLevel );
        
		ht.set_width( currentWidth );
		ht.set_height( currentHeight );
		ht.set_nCh( nChS3 );
        final=clock()-init;
        
        if( displayDebugInfo )
        {
			mexPrintf( "connected a hierarchical template with feature maps %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
			mexEvalString("drawnow");
		}
		
		init = clock();
		
		ht.set_all_deformation_neighborhood( nb );
		ht.set_dict(dicts);
        final=clock()-init;
        
        if( displayDebugInfo )
        {
			mexPrintf( "connected a hierarchical template with dictionaries %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
			mexEvalString("drawnow");
		}
		
		init = clock();
        
		// the following lines display debugging info
		if( displayDebugInfo )
		{
			mexPrintf( "dictionary %d levels\n", dicts.size() );
		    mexPrintf( "dictionary[3] has %d entries\n", dicts[3].size() );
		    mexPrintf( "neighborhood %d levels\n", nb.size() );
		    mexPrintf( "sumMapAllLevel %d levels\n", sumMapAllLevel.size() );
		    mexPrintf( "maxMapAllLevel %d levels\n", maxMapAllLevel.size() );
		    mexPrintf( "argmaxMapAllLevel %d levels\n", argmaxMapAllLevel.size() );
		    mexPrintf( "nChS3 = %d \n", nChS3 );
		    mexPrintf( "%d tuple\n", dicts[3][0]->elements()[0].size() );
		}
        
		ht.FindMyself();
		ht.Expand(2);
		deformedTemplatesLevel1.push_back( ht.ReturnLocsAtLevel(1) );
		deformedTemplatesLevel2.push_back( ht.ReturnLocsAtLevel(2) );
		deformedTemplatesLevel3.push_back( ht.ReturnLocsAtLevel(3) );
		final=clock()-init;
		
		if( displayDebugInfo )
		{
			mexPrintf( "computed the deformed elements for all levels: %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
			mexEvalString("drawnow");
		}
		
	}
}

mxArray* createCellOutput( const vector<SparseT_t>& deformedTemplates )
{
	mwSize dimsOutput[2];
	dimsOutput[0] = deformedTemplates.size(); dimsOutput[1] = 1;
	mxArray* cellarray = mxCreateCellArray( 2, dimsOutput );

    for( int iT = 0; iT < deformedTemplates.size(); ++iT )
    {
		const SparseT_t& st = deformedTemplates[iT];
		dimsOutput[0] = 3; dimsOutput[1] = st.numElement();
        mxArray* pA = mxCreateNumericArray( 2, dimsOutput, mxSINGLE_CLASS, mxREAL );
        /* populate the real part of the created array */
        float* start_of_pr = (float*)mxGetData(pA);
		int cc = 0;
		for( int j = 0; j < dimsOutput[1]; ++j )
		{
			for( int k = 0; k < dimsOutput[0]; ++k )
			{
				start_of_pr[cc] = (float)(st.elements()[j][k]);
				++cc;
			}
		}
        mxSetCell( cellarray, iT, pA );
    }
	return cellarray;
}

mxArray* createCellOutput2( vector<const SparseT_t*> deformedTemplates )
{
	mwSize dimsOutput[2];
	dimsOutput[0] = deformedTemplates.size(); dimsOutput[1] = 1;
	mxArray* cellarray = mxCreateCellArray( 2, dimsOutput );

    for( int iT = 0; iT < deformedTemplates.size(); ++iT )
    {
		SparseT_t st = *(deformedTemplates[iT]);
		dimsOutput[0] = 3; dimsOutput[1] = st.numElement();
        mxArray* pA = mxCreateNumericArray( 2, dimsOutput, mxSINGLE_CLASS, mxREAL );
        /* populate the real part of the created array */
        float* start_of_pr = (float*)mxGetData(pA);
		int cc = 0;
		for( int j = 0; j < dimsOutput[1]; ++j )
		{
			for( int k = 0; k < dimsOutput[0]; ++k )
			{
				start_of_pr[cc] = (float)(st.elements()[j][k]);
				++cc;
			}
		}
        mxSetCell( cellarray, iT, pA );
    }
	return cellarray;
}

/* entry point: load input variables and run the algorithm */
void mexFunction(int nlhs, mxArray *plhs[], 
                 int nrhs, const mxArray *prhs[])   
{
	mxArray* pA;
	const mxArray *cpA;
	mxClassID datatype;
	
    int iT;
	
	/* Following google code style guide, variables of class type should not be
		global or static */
	Dict_t dicts;
	vector<const SparseT_t* > v0(0);
	dicts.push_back( v0 ); // the first level is empty
	Dict_t deformingNeighborhoodAllLevel;
	vector<SparseT_t> deformedTemplatesLevel1;
	vector<SparseT_t> deformedTemplatesLevel2;
	vector<SparseT_t> deformedTemplatesLevel3;

    LOG_INIT_OVERWRITE("main.log");
	LOG_MSG("initialization\n");
	
	/* =============================================
     * Handle input variables.
     * ============================================= 
     */
    /*
	 * input variable 0: paths. 2D cell array.
	 */

	clock_t init, final;
	init=clock();

	cpA = prhs[0];
	paths.clear();
    int numImage = mxGetNumberOfElements(cpA);
	for( int i = 0; i < numImage; ++i )
	{
		const mxArray* cpA2 = mxGetCell(cpA,i);
		int pathDepth = mxGetNumberOfElements(cpA2);
		Path_t onePath;
		for( int j = 0; j < pathDepth; ++j )
		{
            const mxArray* pAString = mxGetCell(cpA2,j);
            unsigned int yLength = mxGetN(pAString)+1;
            char* tmp = (char*)mxCalloc(yLength, sizeof(char));
            mxGetString(pAString,tmp,yLength);
            string a = tmp;
			onePath.push_back( a );
		}
		paths.push_back( onePath );
	}
	
	/*
	 * input variable 1: feature folder path. 1D cell array.
	 */
	cpA = prhs[1];
	int pathDepth = mxGetNumberOfElements(cpA);
	if( !featureFolder.empty() )	featureFolder.clear();
	for( int j = 0; j < pathDepth; ++j )
	{
        const mxArray* pAString = mxGetCell(cpA,j);
        unsigned int yLength = mxGetN(pAString)+1;
        char* tmp = (char*)mxCalloc(yLength, sizeof(char)); //mxCalloc is similar to malloc in C
        mxGetString(pAString,tmp,yLength);
        string a = tmp;
		featureFolder.push_back( a );
	}
	
	
	/*
	 * input variable 2: dictionaries. 2D cell array of SparseTemplate.
	*/
	// inputXXX and destroyXXX should be used in pairs !
	inputDictionaries( prhs[2], dicts ); /* remember to delete dicts */
	
	/*
	 * input variable 3: deformation neighborhood. 2D cell array of SparseTemplate.
	*/
	// Note: In current implementation, there is one deformation neighborhood per template.
	//	The candidate deformations are stored as a list of 3-tuples, i.e. a SparseTemplate.
	//	
	//	Space is wasted if some dimensions of the deformation are independent from one
	//	another. E.g. translation and rotation may be independent.
	//
	//	In future this will be improved. A class will iterate candidate deformations, not
	//	by storing all candidates in a list.
	
	inputDictionaries( prhs[3], deformingNeighborhoodAllLevel ); /* remember to delete dicts */
	final=clock()-init;
	
	if( displayDebugInfo )
	{
		mexPrintf( "loading input: %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
	}
	
	init = clock();
	
	/* Computation */
	compute( dicts, deformingNeighborhoodAllLevel, deformedTemplatesLevel1, deformedTemplatesLevel2, deformedTemplatesLevel3 );
	final=clock()-init;
	
	if( displayDebugInfo )
	{
		mexPrintf( "inference takes: %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
	}
	
	init = clock();
    
	/* Output */
	plhs[0] = createCellOutput( deformedTemplatesLevel1 );
	plhs[1] = createCellOutput( deformedTemplatesLevel2 );
	plhs[2] = createCellOutput( deformedTemplatesLevel3 );
	
	// tmp:
	//plhs[3] = createCellOutput2( dicts[2] );
	//plhs[4] = createCellOutput2( deformingNeighborhoodAllLevel[2] );
	
	final=clock()-init;
	
	if( displayDebugInfo )
	{
		mexPrintf( "output takes: %.3f sec.\n", (double)final / ((double)CLOCKS_PER_SEC) );
	}

	// delete the newed items
	destroyDictionaries( dicts );
	destroyDictionaries( deformingNeighborhoodAllLevel );
	
	LOG_CLOSE;
}
