#include <objrec/nice.h>

#include <iostream>

#include "SemSegContextTree.h"
#include "objrec/baselib/Globals.h"
#include "objrec/baselib/ProgressBar.h"
#include "objrec/baselib/StringTools.h"
#include "objrec/baselib/Globals.h"

#include "objrec/cbaselib/CachedExample.h"
#include "objrec/cbaselib/PascalResults.h"

#include <omp.h>

using namespace OBJREC;
using namespace std;
using namespace NICE;



SemSegContextTree::SemSegContextTree( const Config *conf, const MultiDataset *md )
    : SemanticSegmentation ( conf, &(md->getClassNames("train")) )
{
	string section = "SSContextTree";
	lfcw = new LFColorWeijer(conf);
	
	grid = conf->gI(section, "grid", 10 );
	
	maxSamples = conf->gI(section, "max_samples", 2000 );
	
	minFeats = conf->gI(section, "min_feats", 50 );
	
	maxDepth = conf->gI(section, "max_depth", 20 );
	
	///////////////////////////////////
	// Train Segmentation Context Trees
	//////////////////////////////////

	train ( md );
}

SemSegContextTree::~SemSegContextTree()
{
}

void SemSegContextTree::getBestSplit(const vector<vector<vector<vector<double> > > > &feats, vector<vector<vector<int> > > &currentfeats,const vector<vector<vector<int> > > &labels, int node, int &splitfeat, double &splitval)
{
	int imgCount, featsize;
	try
	{
		imgCount = (int)feats.size();
		featsize = feats[0][0][0].size();
	}
	catch(Exception)
	{
		cerr << "no features computed?" << endl;
	}
	
	double bestig = -numeric_limits< double >::max();
	splitfeat = -1;
	splitval = -1.0;
	
	set<vector<int> >selFeats;
	map<int,int> e;
	int featcounter = 0;
	vector<double> maximum(featsize, -numeric_limits< double >::max());
	vector<double> minimum(featsize, numeric_limits< double >::max());
	
	for(int iCounter = 0; iCounter < imgCount; iCounter++)
	{
		int xsize = (int)currentfeats[iCounter].size();
		int ysize = (int)currentfeats[iCounter][0].size();
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				if(currentfeats[iCounter][x][y] == node)
				{
					featcounter++;
				}
			}
		}
	}
	
	//double fraction = (double)maxSamples/(double)featcounter;
	
	vector<double> fraction(a.size(),0.0);
	for(uint i = 0; i < fraction.size(); i++)
	{
		fraction[i] = ((double)maxSamples)/(a[i]*(double)featcounter*8);
		//cout << "fraction["<<i<<"]: "<< fraction[i] << endl;
	}
	featcounter = 0;
	
	
	for(int iCounter = 0; iCounter < imgCount; iCounter++)
	{
		int xsize = (int)currentfeats[iCounter].size();
		int ysize = (int)currentfeats[iCounter][0].size();
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				if(currentfeats[iCounter][x][y] == node)
				{
					int cn = labels[iCounter][x][y];
					double randD = (double)rand()/(double)RAND_MAX;
					//cout << "class: " << cn << " thres: "<< fraction<<  " rand: " << randD << " scale: " << a[labelmap[cn]] << " newrand: ";
					//randD *= a[labelmap[cn]];
					//cout << randD << endl;
					//getchar();
					if(randD < fraction[labelmap[cn]])
					{
						vector<int> tmp(3,0);
						tmp[0] = iCounter;
						tmp[1] = x;
						tmp[2] = y;
						featcounter++;
						selFeats.insert(tmp);
						
						e[cn] = e[cn]+1;
						for(int f= 0; f < featsize; f++)
						{
							maximum[f] = std::max(maximum[f], feats[iCounter][x][y][f]);
							minimum[f] = std::min(minimum[f], feats[iCounter][x][y][f]);
						}
					}
				}

			}
		}
	}
	
	//cout << "size: " << selFeats.size() << endl;
	
	map<int,int>::iterator mapit;
	double globent = 0.0;
	for ( mapit=e.begin() ; mapit != e.end(); mapit++ )
	{
	  //cout << "class: " << mapit->first << ": " << mapit->second << endl;
		double p = (double)(*mapit).second/(double)featcounter;
		globent += p*log2(p);
	}
	globent = -globent;
	
	if(globent < 0.5)
	{
		cout << "globent to small: " << globent << endl;
		return;
	}
	
	if(featcounter < minFeats)
	{
		cout << "only " << featcounter << " feats in current node -> it's a leaf" << endl;
		return;
	}

	//omp_set_num_threads(2);
#pragma omp parallel for private(mapit)
	for(int f = 0; f < featsize; f++)
	{
		double l_bestig = -numeric_limits< double >::max();
		double l_splitval = -1.0;
	
		set<vector<int> >::iterator it;
		for ( it=selFeats.begin() ; it != selFeats.end(); it++ )
		{
			set<vector<int> >::iterator it2;
			double val = feats[(*it)[0]][(*it)[1]][(*it)[2]] [f];
			
			//cout << "val: " << val << endl;
			if(val == maximum[f] || val == minimum[f])
			{
				continue;
			}
			
			map<int,int> eL, eR;
			int counterL = 0, counterR = 0;
			
			for ( it2=selFeats.begin() ; it2 != selFeats.end(); it2++ )
			{
				int cn = labels[(*it2)[0]][(*it2)[1]][(*it2)[2]];
				if(feats[(*it2)[0]][(*it2)[1]][(*it2)[2]][f] < val)
				{
					//left entropie:
					eL[cn] = eL[cn]+1;
					counterL++;
				}
				else
				{
					//right entropie:
					eR[cn] = eR[cn]+1;
					counterR++;
				}
				
			}
			
			double leftent = 0.0;
			for ( mapit=eL.begin() ; mapit != eL.end(); mapit++ )
			{
				double p = (double)(*mapit).second/(double)counterL;
				leftent += p*log2(p);
			}
			leftent = -leftent;
			
			double rightent = 0.0;
			for ( mapit=eR.begin() ; mapit != eR.end(); mapit++ )
			{
				double p = (double)(*mapit).second/(double)counterR;
				rightent += p*log2(p);
			}
			rightent = -rightent;
			
			double ig = globent - rightent - leftent;
			
			if(ig > l_bestig)
			{
				l_bestig = ig;
				l_splitval = val;
			}
		}
#pragma omp critical
{
		if(l_bestig > bestig)
		{
			bestig = l_bestig;
			splitfeat = f;
			splitval = l_splitval;
		}
}
	}
	//cout << "globent: " << globent <<  " bestig " << bestig << " splitfeat: " << splitfeat << " splitval: " << splitval << endl;
}

void SemSegContextTree::train ( const MultiDataset *md )
{
	const LabeledSet train = * ( *md ) ["train"];
	const LabeledSet *trainp = &train;
	
	ProgressBar pb ( "compute feats" );
	pb.show();
	
	//TODO: Speichefresser!, lohnt sich sparse?
	vector<vector<vector<vector<double> > > > allfeats;
	vector<vector<vector<int> > > currentfeats;
	vector<vector<vector<int> > > labels;
	
	int imgcounter = 0;
	
	LOOP_ALL_S ( *trainp )
	{
		EACH_INFO ( classno,info );

		NICE::ColorImage img;

		std::string currentFile = info.img();

		CachedExample *ce = new CachedExample ( currentFile );

		const LocalizationResult *locResult = info.localization();
		if ( locResult->size() <= 0 )
		{
			fprintf ( stderr, "WARNING: NO ground truth polygons found for %s !\n",
			          currentFile.c_str() );
			continue;
		}

		fprintf ( stderr, "SemSegCsurka: Collecting pixel examples from localization info: %s\n", currentFile.c_str() );

		int xsize, ysize;
		ce->getImageSize ( xsize, ysize );

		vector<vector<int> > tmp = vector<vector<int> >(xsize, vector<int>(ysize,0));
		currentfeats.push_back(tmp);
		labels.push_back(tmp);

		try {
			img = ColorImage(currentFile);
		} catch (Exception) {
			cerr << "SemSeg: error opening image file <" << currentFile << ">" << endl;
			continue;
		}

		Globals::setCurrentImgFN ( currentFile );

		//TODO: resize image?!
		
		vector<vector<vector<double> > > feats;
#if 1
		lfcw->getFeats(img, feats);
#else
		feats = vector<vector<vector<double> > >(xsize,vector<vector<double> >(ysize,vector<double>(3,0.0)));
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				for(int r = 0; r < 3; r++)
				{
					feats[x][y][r] = img.getPixel(x,y,r);
				}
			}
		}
#endif	
		allfeats.push_back(feats);
		
		// getting groundtruth
		NICE::Image pixelLabels (xsize, ysize);
		pixelLabels.set(0);
		locResult->calcLabeledImage ( pixelLabels, ( *classNames ).getBackgroundClass() );
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				classno = pixelLabels.getPixel(x, y);
				labels[imgcounter][x][y] = classno;
				labelcounter[classno]++;
				//if ( forbidden_classes.find ( classno ) != forbidden_classes.end() )
					//continue;
			}
		}
		imgcounter++;
		pb.update ( trainp->count());
		delete ce;
	}
	pb.hide();
	
	map<int,int>::iterator mapit;
	int classes = 0;
	for(mapit = labelcounter.begin(); mapit != labelcounter.end(); mapit++)
	{
		labelmap[mapit->first] = classes;
		labelmapback[classes] = mapit->first;
		classes++;
	}
	
	//balancing
	int featcounter = 0;
	a = vector<double>(classes,0.0);
	for(int iCounter = 0; iCounter < imgcounter; iCounter++)
	{
		int xsize = (int)currentfeats[iCounter].size();
		int ysize = (int)currentfeats[iCounter][0].size();
		for(int x = 0; x < xsize; x++)
		{
			for(int y = 0; y < ysize; y++)
			{
				featcounter++;
				int cn = labels[iCounter][x][y];
				a[labelmap[cn]] ++;
			}
		}
	}
	
	for(int i = 0; i < (int)a.size(); i++)
	{
		a[i] /= (double)featcounter;
		cout << "a["<<i<<"]: " << a[i] << endl;
	}
	
	tree.push_back(Node());
	tree[0].dist = vector<double>(classes,0.0);
	int depth = 0;
	tree[0].depth = depth;
	
	bool allleaf = false;
	while(!allleaf && depth < maxDepth)
	{
		allleaf = true;
		//TODO vielleicht parallel wenn nächste schleife auch noch parallelsiert würde, die hat mehr gewicht
		//#pragma omp parallel for
		int t = (int) tree.size();
		for(int i = 0; i < t; i++)
		{
			if(!tree[i].isleaf && tree[i].left < 0)
			{  
				int splitfeat;
				double splitval;

				getBestSplit(allfeats, currentfeats,labels, i, splitfeat, splitval);
				tree[i].feat = splitfeat;
				tree[i].decision = splitval;
				if(splitfeat >= 0)
				{
					allleaf = false;
					int left = tree.size();
					tree.push_back(Node());
					tree.push_back(Node());
					int right = left+1;
					tree[i].left = left;
					tree[i].right = right;
					tree[left].dist = vector<double>(classes, 0.0);
					tree[right].dist = vector<double>(classes, 0.0);
					tree[left].depth = depth+1;
					tree[right].depth = depth+1;
//#pragma omp parallel for
					for(int iCounter = 0; iCounter < imgcounter; iCounter++)
					{
						int xsize = currentfeats[iCounter].size();
						int ysize = currentfeats[iCounter][0].size();
						for(int x = 0; x < xsize; x++)
						{
							for(int y = 0; y < ysize; y++)
							{
								if(currentfeats[iCounter][x][y] == i)
								{
									if(allfeats[iCounter][x][y][splitfeat] < splitval)
									{ 
										currentfeats[iCounter][x][y] = left;
										tree[left].dist[labelmap[labels[iCounter][x][y]]]++;
									}
									else
									{  
										currentfeats[iCounter][x][y] = right;
										tree[right].dist[labelmap[labels[iCounter][x][y]]]++;
									}
								}
							}
						}
					}
					double lcounter = 0.0, rcounter = 0.0;
					for(uint d = 0; d < tree[left].dist.size(); d++)
					{
						tree[left].dist[d]/=a[d];
						lcounter +=tree[left].dist[d];
						tree[right].dist[d]/=a[d];
						rcounter +=tree[right].dist[d];
					}
					for(uint d = 0; d < tree[left].dist.size(); d++)
					{
						tree[left].dist[d]/=lcounter;
						tree[right].dist[d]/=rcounter;
					}
				}
				else
				{
					tree[i].isleaf = true;
				}
				//TODO: probability ermitteln
			}
		}
		
		//TODO: features neu berechnen!
		depth++;
		cout << "d: " << depth << endl;
	}

}

void SemSegContextTree::semanticseg ( CachedExample *ce, NICE::Image & segresult,GenericImage<double> & probabilities )
{
	int xsize;
	int ysize;
	ce->getImageSize ( xsize, ysize );
	
	int numClasses = classNames->numClasses();
	
	fprintf (stderr, "ContextTree classification !\n");

	probabilities.reInit ( xsize, ysize, numClasses, true );
	probabilities.setAll ( 0 );

	NICE::ColorImage img;

	std::string currentFile = Globals::getCurrentImgFN();
	
	try {
		img = ColorImage(currentFile);
	} catch (Exception) {
		cerr << "SemSeg: error opening image file <" << currentFile << ">" << endl;
		return;
	}
	
	//TODO: resize image?!
		
	vector<vector<vector<double> > > feats;
	
#if 1
	lfcw->getFeats(img, feats);
#else
	feats = vector<vector<vector<double> > >(xsize,vector<vector<double> >(ysize,vector<double>(3,0.0)));
	for(int x = 0; x < xsize; x++)
	{
		for(int y = 0; y < ysize; y++)
		{
			for(int r = 0; r < 3; r++)
			{
				feats[x][y][r] = img.getPixel(x,y,r);
			}
		}
	}
#endif
	
	bool allleaf = false;
	
	vector<vector<int> > currentfeats = vector<vector<int> >(xsize, vector<int>(ysize,0));
	int depth = 0;
	while(!allleaf)
	{
		allleaf = true;
		//TODO vielleicht parallel wenn nächste schleife auch noch parallelsiert würde, die hat mehr gewicht
		//#pragma omp parallel for
		int t = (int) tree.size();
		for(int i = 0; i < t; i++)
		{
			for(int x = 0; x < xsize; x++)
			{
				for(int y = 0; y < ysize; y++)
				{
					int t = currentfeats[x][y];
					if(tree[t].left > 0)
					{
						allleaf = false;
						if(feats[x][y][tree[t].feat] < tree[t].decision)
						{
							currentfeats[x][y] = tree[t].left;
						}
						else
						{
							currentfeats[x][y] = tree[t].right;
						}
					}
				}
			}
		}
		
		//TODO: features neu berechnen! analog zum training
		
		depth++;
	}
	
	//finales labeln:
	 long int offset = 0;
	for(int x = 0; x < xsize; x++)
	{
		for(int y = 0; y < ysize; y++,offset++)
		{
			int t = currentfeats[x][y];
			double maxvalue = - numeric_limits<double>::max(); //TODO: das muss nur pro knoten gemacht werden, nicht pro pixel
			int maxindex = 0;
			for(uint i = 0; i < tree[i].dist.size(); i++)
			{
				probabilities.data[labelmapback[i]][offset] = tree[t].dist[i];
				if(tree[t].dist[i] > maxvalue)
				{
					maxvalue = tree[t].dist[i];
					maxindex = labelmapback[i];
				}
				segresult.setPixel(x,y,maxindex);
			}
		}
	}
}