/* 
	Hierarchical localization

It is a light-weight tree class. 

Usage:
	HierTemplate h(0,0,0,3); // S3 localization
	h.set_width(...);
	h.set_height(...);
	h.set_nChn(...);
	h.set_all_sum_map(...);
	h.set_all_max_map(...);
	h.set_all_argmax_map(...);
	h.set_dict(..);
	h.set_all_deformation_neighborhood(...);

	h.FindMyself();
	h.Expand(2);
	vector<SparseT_t> gabors = h.ReturnLocsAtLevel(1);
	vector<SparseT_t> shapemotifs = h.ReturnLocsAtLevel(2);
	vector<SparseT_t> deformed_shapescript = h.ReturnLocsAtLevel(3);
*/

#include <vector>
#include "assert.h"
#include "Log.h"
#include <string.h>
#include "mex.h"    // TODO: make it optional
using namespace std;

#define NEGMAX -1e-10

// a simple flat sparse template: a list of 3-tuples
// TODO: add assertions for 3-tuple checking
template<class Param_t>
class SparseTemplate
{
	typedef vector< vector<Param_t> > V_t;
	private:
		V_t elements_;
	public:
		SparseTemplate() { V_t in; elements_ = in; }
		SparseTemplate( const SparseTemplate &in ) { elements_ = in.elements(); }
		SparseTemplate( V_t in ) { elements_ = in; }

		V_t elements() const {return elements_;}
		void set_elements( V_t in ) { elements_ = in; }
		
		int numElement() const { return elements_.size(); }
};

///

// TODO: the current interface is not clean enough.
/*
	HierTemplate should only contain:
	1. its own free parameters: row,col,chn.
	2. links to its children.
	3. link to the dictionaries of flat templates, and related info
	4. link to the data
*/

template <class Pixel_t, class Param_t>
class HierTemplate : public SparseTemplate<Param_t> {

	public:
		typedef vector< vector< const SparseTemplate<Param_t>* > > Dictionaries_t; // a 2D array of sparse templates
		typedef SparseTemplate<Param_t> SparseT_t;

	private:
	
		Log* log_;	// to output debugging info into file
		int level_;	// level of visual dictionary
		Param_t row_, col_, chn_; /* canonical row, column, channel */
		int width_, height_, nCh_;
		
		// ? no need to use HierTemplate*, since children nodes will not be shared. So it does not save space.
		vector< HierTemplate > children_;	/* children nodes */
		
		/* dictionary of templates: access by dict_[level_][chn_] */
        // TODO: summarize this into DictionaryManager
		Dictionaries_t dict_;

		/* collection of sum/max/argmax maps at all levels */
		// methods in HierTemplate do not modify feature maps
        
        // TODO: summarize the feature maps into another class FeatureManager
        //      Then only need to pass one pointer to children.
		vector<const Pixel_t**> all_sum_map_;
		vector<const Pixel_t**> all_max_map_;
		vector<const int**> all_argmax_map_;
		
		/* mapping from pixel value of argmax_map_ to row, col, chn */
		Dictionaries_t all_deformation_neighborhood_;
		
		/* The dimension of the following maps are #channel by (#row * #col) */
		// const Pixel_t **sum_map_;		/* the current SUM map it is visiting */
		// const Pixel_t **max_map_;
		// const int **argmax_map_;
		// but it can update the inhibition map of the corresponding level
		bool **inhibition_map_; // inhibition_map_ should be deleted outside this class
		bool translated_;

	public:
	
		HierTemplate() {
			row_=0; col_=0; level_=0;
			nCh_ = 0;
			translated_ = false;
		};
	
		HierTemplate ( Param_t row, Param_t col, Param_t chn, int level ){ 
			row_=row; col_=col; chn_=chn; level_=level;
			nCh_ = 0; width_ = 0; height_ = 0;
			translated_ = false;
		};

		/* The copy constructor copies localization. */
		HierTemplate(const HierTemplate& in) {
			row_ = in.row();
			col_ = in.col();
			chn_ = in.chn();
			level_ = in.level();
			translated_ = in.translated();
            children_ = in.children();
            this->set_all_sum_map( in.all_sum_map() );
            this->set_all_max_map( in.all_max_map() );
            this->set_all_argmax_map( in.all_argmax_map() );

            this->set_width( in.width() ); // assume Sx maps share the same lattice size 
            this->set_height( in.height() );

            //
            this->set_all_deformation_neighborhood( in.all_deformation_neighborhood() );
            this->set_dict( in.dict() );
		}
		
		~HierTemplate() {};
		
		/* Accessors and Mutators */
        bool translated() const { return translated_; }
		int level() const { return level_; }
		void set_level(int in) { level_ = in; }
		int width() const { return width_; }
		void set_width(int in) { width_ = in; }
		int height() const { return height_; }
		void set_height(int in) { height_ = in; }
		int nCh() const { return nCh_; }
		void set_nCh(int in) { nCh_ = in; }
		Param_t row() const { return row_; }
		void set_row(Param_t in) { row_ = in; }
		Param_t col() const { return col_; }
		void set_col(Param_t in) { col_ = in; }
		Param_t chn() const { return chn_; }
		void set_chn(Param_t in) { chn_ = in; }
		vector<HierTemplate> children() const { return children_; }
		void set_children(vector<HierTemplate> in) { children_ = in; }
		bool** inhibition_map() const { return inhibition_map_; }
		void set_inhibition_map(bool** in) { inhibition_map_ = in; }
		vector<const Pixel_t**> all_sum_map() const { return all_sum_map_; }
		void set_all_sum_map(vector<const Pixel_t**> in) { all_sum_map_ = in; }
		vector<const Pixel_t**> all_max_map() const { return all_max_map_; }
		void set_all_max_map(vector<const Pixel_t**> in) { all_max_map_ = in; }
		vector<const int**> all_argmax_map() const { return all_argmax_map_; }
		void set_all_argmax_map(vector<const int**> in) { all_argmax_map_ = in; }
        Dictionaries_t all_deformation_neighborhood() const { return all_deformation_neighborhood_; }
		void set_all_deformation_neighborhood( Dictionaries_t in ) { all_deformation_neighborhood_ = in; }
		void set_log(Log* in) { log_ = in; }
        Dictionaries_t dict() const { return dict_; }
		void set_dict(Dictionaries_t in) { dict_ = in; }

		/* ==================== regular functions ========================= */
		
		/* compute the top (row_, col_, ind_) with largest pixel value from sum_map_ and optionally inhibition_map_ */
		void FindNextBest(bool useInhibition = false, bool inhibitiOthers = false)
		{
            /*mexPrintf("inside FindNextBest()\n");
            mexPrintf("%d %d\n", level_, all_sum_map_.size() );
            mexPrintf("%d %d\n", width_, height_ );*/
            
			const Pixel_t **sum_map_ = all_sum_map_[level_];
            nCh_ = dict_[level_].size();
			
			//assert( sum_map_ != NULL );
            if( sum_map_ == NULL )
            {
                mexErrMsgTxt( "sum map is NULL" );
            }
			
			int bestCh = -1;
			int bestRow = -1;
			int bestCol = -1;
			
			if( useInhibition )
			{
				// check whether inhibition map is set
				assert( inhibition_map_ != NULL );
				Pixel_t maxVal = NEGMAX;
				
				for( int ch = 0; ch < nCh_; ++ch )
				{
					int j = 0;
					for( int c = 0; c < width_; ++c )
					for( int r = 0; r < height_; ++r )
					{
						Pixel_t val = sum_map_[ch][j];
						if( val > maxVal && inhibition_map_[ch][j])
						{
							maxVal = val;
							bestCh = ch;
							bestRow = r;
							bestCol = c;
						}
						++j;
					}
				}
			}
			else
			{
                // mexPrintf("inhibition is not enabled.\n");
				Pixel_t maxVal = NEGMAX;
				for( int ch = 0; ch < nCh_; ++ch )
				{
					int j = 0;
					for( int c = 0; c < width_; ++c )
					for( int r = 0; r < height_; ++r )
					{
						Pixel_t val = sum_map_[ch][j];
						if( val > maxVal )
						{
							maxVal = val;
							bestCh = ch;
							bestRow = r;
							bestCol = c;
						}
						++j;
					}
				}
			}
			// mexPrintf("found myself\n");
            /*if( level_ == 3)
            {
                mexPrintf("currently we don't search globally for S3 center\n");
                bestRow = (int)height_/2;
                bestCol = (int)width_/2;
                bestCh = 0;
            }*/
            
			row_ = bestRow;
			col_ = bestCol;
			chn_ = bestCh;
            translated_ = true;
            // mexPrintf("top node found at (%.1f %.1f %.1f)\n", (float)row_, (float)col_, (float)chn_);
		}
		/* below: duplicate methods */
		void FindMyself(bool useInhibition = false, bool inhibitiOthers = false){FindNextBest(useInhibition,inhibitiOthers);}

		/* compute the top (row_, col_, ind_) from argmax_map_ and optionally inhibition_map_ */
		void FindMyselfNearBy(bool inhibitionIsOn = false)
		{
			nCh_ = dict_[level_].size();
			// remind that the parent location should be found first
            if( !translated_ )
            {
                mexErrMsgTxt("should be translated");
            }
			
			const int** argmax_map = all_argmax_map_[level_];
			//assert( argmax_map != NULL );
			
			if( inhibitionIsOn )
			{
				assert( inhibition_map_ != NULL );
				// TODO: fill in argmax operation when inhibition is on
			}
			else
			{
				Param_t dr, dc, destc;
				bool valid = true;
				if( chn_ < 0 || chn_ >= nCh_ )	valid = false;
				if( row_ < 0 || row_ >= height_ )	valid = false;
				if( col_ < 0 || col_ >= width_ )	valid = false;
				if( valid )
				{
					int ind = argmax_map[(int)chn_][(int)(row_+col_ * height_)];
					// TODO: implement a class with method that retrieves the stored candidate deformations
					//		const Param_t[3] getDeform(ind,this_Chn,(optional)parent_Chn);
					const SparseT_t* neighborhood = all_deformation_neighborhood_[level_][chn_];
					dr = neighborhood->elements()[ind][0];
					dc = neighborhood->elements()[ind][1];
					destc = neighborhood->elements()[ind][2];

					// translate
					row_ += dr;
					col_ += dc;
					chn_ = destc;
                    translated_ = true;
				}
				else
				{
					// mexPrintf("not valid!! %d, %d, %d, width = %d, height = %d, level = %d, nCh = %d \n", (int)row_, (int)col_, (int)chn_, (int)width_, (int)height_, (int)level_, (int)nCh_ );
				}
				
			}
		}
		
		/* translate the template: itself and its children */
		void Translate(Param_t dr, Param_t dc) 
		{
			row_ = row_ + dr;
			col_ = col_ + dc;
			/* recursion: */
			for( int i = 0; i < children_.size(); ++i )
			{
				children_[i].Translate( dr, dc );
			}
			translated_ = true;
		}
		
		/* find the tree of hierarchical localization */
		void Expand( int depth = 0 ) 
		{
            
			if( depth == 0 )
			{
				// do nothing
			}
			else
			{
				/* find the matching template at canonical position */
				const SparseT_t *p = dict_[level_][chn_];

				//assert( p != NULL );
                if( p == NULL )
                {
                    mexErrMsgTxt( "cannot find myself in the dictionary" );
                }
				
				/* copy a translated version of children nodes from the dictionary template */
				children_.clear();

                int count = 1;
				for( int i = 0; i < p->elements().size(); ++i )
				{
					vector<Param_t> tmp = p->elements()[i];
                    ++count;
					HierTemplate a( tmp[0], tmp[1], tmp[2], level_-1 );
					a.set_width( this->width() ); // assume Sx maps share the same lattice size 
                    a.set_height( this->height() );
					
					a.Translate( row_, col_ );
                    a.set_all_sum_map( this->all_sum_map() );
                    a.set_all_max_map( this->all_max_map() );
                    a.set_all_argmax_map( this->all_argmax_map() );

                    //
                    a.set_all_deformation_neighborhood( this->all_deformation_neighborhood() );
                    a.set_dict( this->dict() );
					children_.push_back( a );
				}


				/* local movement of children nodes */
				for( int i = 0; i < children_.size(); ++i )
				{
                    //mexPrintf("iii child at %.3f, %.3f\n",(float)i -> row(), (float)i->col());
					children_[i].FindMyselfNearBy();
				}
				
				/* recursion: Do the same for children nodes */
				if( depth > 1 )
				{
					for( int i = 0; i < children_.size(); ++i )
					{
						children_[i].Expand( depth - 1 );
					}
				}
			}
		}
		
		/* return a list of (row,column,channel) at the specific tree level (0 is leaf) */
		SparseTemplate<Param_t> ReturnLocsAtLevel( int level ) const
		{
			typedef vector< vector< Param_t > > Vec2D_t;
			if( level_ < level )
			{
				Vec2D_t a(0);
				return a;
			}
			else if( level_ == level )
			{
				Vec2D_t a;
				vector<Param_t> aa;
				aa.push_back( row_ ); aa.push_back( col_ ); aa.push_back( chn_ );
                //mexPrintf("at level %d, pushed %d, %d, %d\n", level_, (int)row_, (int)col_, (int)chn_ );
				a.push_back(aa);
				return a;
			}
			else
			{
				/* recursion : */
               // mexPrintf("recursion\n" );
				Vec2D_t a;
				int cc = 0;
				for( int i = 0; i < children_.size(); ++i )
				{
					SparseTemplate<Param_t> aa = children_[i].ReturnLocsAtLevel( level );
					/* append aa to the end of a */
					
					for( int j = 0; j < aa.elements().size(); ++j )
					{
						a.push_back( aa.elements()[j] );
					}
					++cc;
				}
				return a;
			}
		}
};
