// Beispielhafter Aufruf: BUILD_x86_64/progs/testSemanticSegmentation -config <CONFIGFILE>

/**
* @file testSemanticSegmentation.cpp
* @brief test semantic segmentation routines
* @author Erik Rodner
* @date 03/20/2008
*/

#ifdef NICE_USELIB_OPENMP
#include <omp.h>
#endif

#include <objrec/baselib/Config.h>
#include <objrec/baselib/StringTools.h>
#include <objrec/baselib/ICETools.h>

#include <objrec/semanticsegmentation/SemanticSegmentation.h>
#include <objrec/semanticsegmentation/SemSegLocal.h>
#include <objrec/semanticsegmentation/SemSegSTF.h>
#include <objrec/semanticsegmentation/SemSegCsurka.h>
#include <objrec/semanticsegmentation/SemSegRegionBased.h>

#include <fstream>

using namespace OBJREC;

using namespace NICE;

using namespace std;

void updateMatrix(const NICE::Image & img, const NICE::Image & gt,
																		NICE::Matrix & M, const set<int> & forbidden_classes)
{
	double subsamplex = gt.width() / (double)img.width();
	double subsampley = gt.height() / (double)img.height();

	for (int y = 0 ; y < gt.height() ; y++)
		for (int x = 0 ; x < gt.width() ; x++)
		{
			int xx = (int)(x / subsamplex);
			int yy = (int)(y / subsampley);

			if (xx < 0) xx = 0;

			if (yy < 0) yy = 0;

			if (xx > img.width() - 1) xx = img.width() - 1;

			if (yy > img.height() - 1) yy = img.height() - 1;

			int cimg = img.getPixel(xx, yy);

			int gimg = gt.getPixel(x, y);

			if (forbidden_classes.find(gimg) == forbidden_classes.end())
			{
				M(gimg, cimg)++;
			}
		}
}

/**
 test semantic segmentation routines
*/
int main(int argc, char **argv)
{
	std::set_terminate(__gnu_cxx::__verbose_terminate_handler);

	Config conf(argc, argv);

	bool show_result = conf.gB("debug", "show_results", false);

	bool write_results = conf.gB("debug", "write_results", false);

	bool write_results_pascal = conf.gB("debug", "write_results_pascal", false);

	std::string resultdir = conf.gS("debug", "resultdir", ".");

	if (write_results)
	{
		cerr << "Writing Results to " << resultdir << endl;
	}

	MultiDataset md(&conf);

	const ClassNames & classNames = md.getClassNames("train");

	//SemanticSegmentation *semseg = new SemSegLocal ( &conf, &md );
	//SemanticSegmentation *semseg = new SemSegSTF ( &conf, &md );
	//SemanticSegmentation *semseg = new SemSegCsurka ( &conf, &md);

	SemanticSegmentation *semseg = new SemSegRegionBased(&conf, &md);

	const LabeledSet *testFiles = md["test"];
	NICE::Matrix M(classNames.getMaxClassno() + 1, classNames.getMaxClassno() + 1);
	M.set(0);

	set<int> forbidden_classes;

	std::string forbidden_classes_s = conf.gS("analysis", "forbidden_classes", "");

	classNames.getSelection(forbidden_classes_s, forbidden_classes);

	ProgressBar pb("Semantic Segmentation Analysis");
	pb.show();

	int fileno = 0;

	LOOP_ALL_S(*testFiles)
	{
		EACH_INFO(classno, info);
		pb.update(testFiles->count());

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

		NICE::Image lm;
		GenericImage<double> probabilities;
		
		if (info.hasLocalizationInfo())
		{
			const LocalizationResult *l_gt = info.localization();

			lm.resize(l_gt->xsize, l_gt->ysize);
			lm.set(0);
			l_gt->calcLabeledImage(lm, classNames.getBackgroundClass());
		}

		semseg->semanticseg(file, lm, probabilities);
		fprintf(stderr, "testSemanticSegmentation: Segmentation finished !\n");

		NICE::Image lm_gt;

		if (info.hasLocalizationInfo())
		{
			const LocalizationResult *l_gt = info.localization();

			lm_gt.resize(l_gt->xsize, l_gt->ysize);
			lm_gt.set(0);

			fprintf(stderr, "testSemanticSegmentation: Generating Labeled NICE::Image (Ground-Truth)\n");
			l_gt->calcLabeledImage(lm_gt, classNames.getBackgroundClass());
		}

		std::string fname = StringTools::baseName(file, false);

		if (write_results_pascal)
		{
			NICE::Image pascal_lm(lm.width(), lm.height());
			int backgroundClass = classNames.getBackgroundClass();

			for (int y = 0 ; y < lm.height(); y++)
				for (int x = 0 ; x < lm.width(); x++)
				{
					int v = lm.getPixel(x, y);

					if (v == backgroundClass)
						pascal_lm.setPixel(x, y, 255);
					else
						pascal_lm.setPixel(x, y, 255 - v - 1);
				}

			char filename[1024];

			char *format = (char *)"pgm";
			sprintf(filename, "%s/%s.%s", resultdir.c_str(), fname.c_str(), format);

			pascal_lm.write(filename);
		}

		if (show_result || write_results)
		{
			NICE::ColorImage orig(file);
			NICE::ColorImage rgb;
			NICE::ColorImage rgb_gt;
			classNames.labelToRGB(lm, rgb);
			classNames.labelToRGB(lm_gt, rgb_gt);

			if (write_results)
			{
				char filename[1024];
				char *format = "ppm";
				sprintf(filename, "%06d.%s", fileno, format);
				std::string origfilename = resultdir + "/orig_" + string(filename);
				cerr << "Writing to file " << origfilename << endl;
				orig.write(origfilename);
				rgb.write(resultdir + "/result_" + string(filename));
				rgb_gt.write(resultdir + "/groundtruth_" + string(filename));
			}

			if (show_result)
			{
#ifndef NOVISUAL
				showImage(rgb, "Result");
				showImage(rgb_gt, "Groundtruth");
				showImage(orig, "Input");
#endif
			}
		}

//#pragma omp critical
		updateMatrix(lm, lm_gt, M, forbidden_classes);

		cerr << M << endl;
		fileno++;
	}

	pb.hide();

	double overall = 0.0;
	double sumall = 0.0;
	for(int r = 0; r < (int)M.rows(); r++)
	{
		for(int c = 0; c < (int)M.cols(); c++)
		{
			if(r == c)
				overall += M(r,c);
			sumall += M(r,c);
		}
	}
	
	overall /= sumall;
	
	// normalizing M using rows
	for (int r = 0 ; r < (int)M.rows() ; r++)
	{
		double sum = 0.0;

		for (int c = 0 ; c < (int)M.cols() ; c++)
			sum += M(r, c);

		if (fabs(sum) > 1e-4)
			for (int c = 0 ; c < (int)M.cols() ; c++)
				M(r, c) /= sum;
	}

	cerr << M << endl;
	
	double avg_perf = 0.0;
	int classes_trained = 0;

	for (int r = 0 ; r < (int)M.rows() ; r++)
	{
		if ((classNames.existsClassno(r)) && (forbidden_classes.find(r) == forbidden_classes.end()))
		{
			avg_perf += M(r, r);
			classes_trained++;
		}
	}

	if (write_results)
	{
		ofstream fout((resultdir + "/res.txt").c_str(), ios::out);
		fout <<  "overall: " << overall << endl;
		fout << "Average Performance " << avg_perf / (classes_trained) << endl;
		fout << "Lower Bound " << 1.0  / classes_trained << endl;

		for (int r = 0 ; r < (int)M.rows() ; r++)
		{
			if ((classNames.existsClassno(r)) && (forbidden_classes.find(r) == forbidden_classes.end()))
			{
				std::string classname = classNames.text(r);
				fout << classname.c_str() << ": " << M(r, r) << endl;
			}
		}

		fout.close();
	}

	cerr <<  "overall: " << overall << endl;
	fprintf(stderr, "Average Performance %f\n", avg_perf / (classes_trained));
	//fprintf(stderr, "Lower Bound %f\n", 1.0 / classes_trained);
	for (int r = 0 ; r < (int)M.rows() ; r++)
	{
		if ((classNames.existsClassno(r)) && (forbidden_classes.find(r) == forbidden_classes.end()))
		{
			std::string classname = classNames.text(r);
			fprintf(stderr, "%s: %f\n", classname.c_str(), M(r, r));
		}
	}

	delete semseg;

	return 0;
}