/** 
* @file SemSegContextTree.h
* @brief Context Trees -> Combination of decision tree and context information
* @author Björn Fröhlich
* @date 29.11.2011

*/
#ifndef SemSegContextTreeINCLUDE
#define SemSegContextTreeINCLUDE

#include "SemanticSegmentation.h"
#include <core/vector/VVector.h>
#include "vislearning/features/localfeatures/LFColorWeijer.h"
#include "objrec/segmentation/RegionSegmentationMethod.h"

namespace OBJREC {
	
class Operation;

class TreeNode
{
public:
	/** probabilities for each class */
	std::vector<double> probs;
	
	/** left child node */
	int left;
	
	/** right child node */
	int right;
	
	/** position of feat for decision */
	Operation *feat;
	
	/** decision stamp */
	double decision;
	
	/** is the node a leaf or not */
	bool isleaf;
	
	/** distribution in current node */
	std::vector<double> dist;
	
	/** depth of the node in the tree */
	int depth;
	
	/** simple constructor */
	TreeNode():left(-1),right(-1),feat(NULL), decision(-1.0), isleaf(false){}
	
	/** standard constructor */
	TreeNode(int _left, int _right, Operation *_feat, double _decision):left(_left),right(_right),feat(_feat), decision(_decision),isleaf(false){}
};
	
struct Features{
	NICE::MultiChannelImageT<double> *feats;
	MultiChannelImageT<int> *cfeats;
	int cTree;
	std::vector<TreeNode> *tree;
	NICE::MultiChannelImageT<double> *integralImg;
};

class ValueAccess
{
public:
	virtual double getVal(const Features &feats, const int &x, const int &y, const int &channel) = 0;
	virtual std::string writeInfos() = 0;
};

class Operation
{
protected:
	int x1, y1, x2, y2, channel1, channel2;
	ValueAccess *values;
public:
	virtual void set(int _x1, int _y1, int _x2, int _y2, int _channel1, int _channel2, ValueAccess *_values)
	{
		x1 = _x1;
		y1 = _y1;
		x2 = _x2;
		y2 = _y2;
		channel1 = _channel1;
		channel2 = _channel2;
		values = _values;
	}
	/**
	 * @brief abstract interface for feature computation
	 * @param feats features
	 * @param cfeats number of tree node for each pixel
	 * @param tree current tree
	 * @param x current x position
	 * @param y current y position
	 * @return double distance
	 **/
	virtual double getVal(const Features &feats, const int &x, const int &y) = 0;
	virtual Operation* clone() = 0;
	virtual std::string writeInfos() = 0;
	
	inline void getXY(const Features &feats, int &xsize, int &ysize)
	{
		xsize = feats.feats->width();
		ysize = feats.feats->height();
	}
};
  
/** Localization system */
class SemSegContextTree : public SemanticSegmentation
{
	/** Segmentation Method */
	RegionSegmentationMethod *segmentation;
			
	/** tree -> saved as vector of nodes */
	std::vector<std::vector<TreeNode> > forest;
	
	/** local features */
	LFColorWeijer *lfcw;
	
	/** number of featuretype -> currently: local and context features = 2 */
	int ftypes;
	
	/** distance between features */
	int grid;
	
	/** maximum samples for tree  */
	int maxSamples;
	
	/** size for neighbourhood */
	int windowSize;
	
	/** how many feats should be considered for a split */
	int featsPerSplit;
	
	/** count samples per label */
	std::map<int,int> labelcounter;
	
	/** map of labels */
	std::map<int,int> labelmap;
	
	/** map of labels inverse*/
	std::map<int,int> labelmapback;
	
	/** scalefactor for balancing for each class */
	std::vector<double> a;
	
	/** the minimum number of features allowed in a leaf */
	int minFeats;
	
	/** maximal depth of tree */
	int maxDepth;
	
	/** operations for pairwise features */
	std::vector<Operation*> ops;
	
	/** operations for pairwise context features */
	std::vector<Operation*> cops;
	
	std::vector<ValueAccess*> calcVal;
	
	/** vector of all possible features */
	std::vector<Operation*> featsel;
	
	/** use alternative calculation for information gain */
	bool useShannonEntropy;
	
	/** Classnames */
	ClassNames classnames;

	/** train selection */
	std::set<int> forbidden_classes;
	
	/** Configfile */
	const Config *conf;
	
	/** use pixelwise labeling or regionlabeling with additional segmenation */
	bool pixelWiseLabeling;
	
	/** use Gaussian distributed features based on the feature position */
	bool useGaussian;
	
	/** Number of trees used for the forest */
	int nbTrees;
	
    public:
	/** simple constructor */
	SemSegContextTree( const NICE::Config *conf, const MultiDataset *md );
      
	/** simple destructor */
	virtual ~SemSegContextTree();

	/**
	 * test a single image
	 * @param ce input data
	 * @param segresult segmentation results
	 * @param probabilities probabilities for each pixel
	 */
	void semanticseg ( CachedExample *ce,   NICE::Image & segresult,  NICE::MultiChannelImageT<double> & probabilities );
	
	/**
	 * the main training method
	 * @param md training data 
	 */
	void train ( const MultiDataset *md );
	
	
	/**
	 * @brief computes integral image of given feats
	 *
	 * @param currentfeats input features
	 * @param integralImage output image (must be initilized)
	 * @return void
	 **/
	void computeIntegralImage(const NICE::MultiChannelImageT<int> &currentfeats, const NICE::MultiChannelImageT<int> &lfeats, NICE::MultiChannelImageT<double> &integralImage);
	
	/**
	 * compute best split for current settings
	 * @param feats features
	 * @param currentfeats matrix with current node for each feature
	 * @param labels labels for each feature
	 * @param node current node
	 * @param splitfeat output feature position
	 * @param splitval 
	 * @return best information gain
	 */
	double getBestSplit(std::vector<NICE::MultiChannelImageT<double> > &feats, std::vector<NICE::MultiChannelImageT<int> > &currentfeats, std::vector<NICE::MultiChannelImageT<double> > &integralImgs, const std::vector<NICE::MatrixT<int> > &labels, int node, Operation *&splitop, double &splitval, const int &tree);
	
	/**
	 * @brief computes the mean probability for a given class over all trees
	 * @param x x position
	 * @param y y position
	 * @param channel current class
	 * @param currentfeats information about the nodes
	 * @return double mean value
	 **/
	inline double getMeanProb(const int &x,const int &y,const int &channel, const MultiChannelImageT<int> &currentfeats);

};


} // namespace

#endif