// 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 "core/basics/Config.h"
#include "core/basics/StringTools.h"
#include <vislearning/baselib/ICETools.h>

#include <semseg3d/image/MultiChannelImage3DT.h>
#include <semseg3d/semseg/SemanticSegmentation.h>
#include <semseg3d/semseg/SemSegContextTree.h>

#include <core/basics/ResourceStatistics.h>

#include <fstream>
#include <vector>

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 );
  
  ResourceStatistics rs;
  
  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 = NULL;
  semseg = new SemSegContextTree( &conf, &md );
  const LabeledSet *testFiles = md["test"];

  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, imageno = 0;

	vector< int > zsizeVec;
	semseg->getDepthVector( testFiles, zsizeVec );

	int depthCount = 0, idx = 0;
	vector< string > filelist;
	NICE::MultiChannelImageT<double> segresult;
	NICE::MultiChannelImageT<double> gt;
	std::vector< NICE::Matrix > M_vec;

	LOOP_ALL_S( *testFiles )
  {
    EACH_INFO( classno, info );
    std::string file = info.img();
    filelist.push_back( file );
		depthCount++;

		NICE::Image lm;
		NICE::Image lm_gt;
		if ( info.hasLocalizationInfo() )
    {
      const LocalizationResult *l_gt = info.localization();

      lm.resize( l_gt->xsize, l_gt->ysize );
      lm.set( 0 );
			
			lm_gt.resize( l_gt->xsize, l_gt->ysize );
      lm_gt.set( 0 );
      
			l_gt->calcLabeledImage( lm, classNames.getBackgroundClass() );
			fprintf( stderr, "testSemanticSegmentation: Generating Labeled NICE::Image (Ground-Truth)\n" );
      l_gt->calcLabeledImage( lm_gt, classNames.getBackgroundClass() );
    }
    segresult.addChannel( lm );
    gt.addChannel( lm_gt );

		if ( depthCount < zsizeVec[idx] ) continue;

    NICE::MultiChannelImage3DT<double> probabilities;
		NICE::MultiChannelImage3DT<double> imgData;
		semseg->make3DImage( filelist, imgData );
		semseg->semanticseg( imgData, segresult, probabilities, filelist );

    fprintf( stderr, "testSemanticSegmentation: Segmentation finished !\n" );

    // save to file
		for (int z = 0; z < zsizeVec[idx]; z++)
		{
			std::string fname = StringTools::baseName( filelist[z], false );

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

				for ( int y = 0 ; y < segresult.height(); y++ )
				{
					for ( int x = 0 ; x < segresult.width(); x++ )
					{
						int v = segresult.get( x, y, (uint)z );

						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( filelist[z] );
				NICE::ColorImage rgb;
				NICE::ColorImage rgb_gt;

				NICE::Image lm( segresult.width(), segresult.height() );
				NICE::Image lm_gt( segresult.width(), segresult.height() );
				for ( int y = 0 ; y < segresult.height(); y++ )
				{
					for ( int x = 0 ; x < segresult.width(); x++ )
					{
						lm.setPixel( x, y, segresult.get( x, y, (uint)z ) );
						lm_gt.setPixel( x, y, gt.get( x, y, (uint)z ) );
					}
				}
				classNames.labelToRGB( lm, rgb );
				classNames.labelToRGB( lm_gt, rgb_gt );

				if ( write_results )
				{
					char filename[1024];
					char *format = ( char * )"ppm";
					sprintf( filename, "%03d_%03d.%s", imageno, 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 ) );
					fileno++;
				}

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

//#pragma omp critical
		for (int z = 0; z < zsizeVec[idx]; z++)
		{
			NICE::Image lm( segresult.width(), segresult.height() );
			NICE::Image lm_gt( segresult.width(), segresult.height() );
			for ( int y = 0 ; y < segresult.height(); y++ )
			{
				for ( int x = 0 ; x < segresult.width(); x++ )
				{
					lm.setPixel( x, y, segresult.get( x, y, (uint)z ) );
					lm_gt.setPixel( x, y, gt.get( x, y, (uint)z ) );
				}
			}
			NICE::Matrix M( classNames.getMaxClassno() + 1, classNames.getMaxClassno() + 1 );
			M.set( 0 );
			updateMatrix( lm, lm_gt, M, forbidden_classes );
			M_vec.push_back( M );
			
			cerr << M << endl;
		}

		// prepare for new 3d image
		filelist.clear();
		NICE::MultiChannelImageT<double> segresult;
		NICE::MultiChannelImageT<double> gt;
		depthCount = 0;
		idx++;
		imageno++;

		pb.update( testFiles->count() );
  }
	segresult.freeData();
  pb.hide();

  long maxMemory;
  rs.getMaximumMemory(maxMemory);
  cerr << "Maximum memory used: " << maxMemory << " KB" << endl;

  double overall = 0.0;
  double sumall = 0.0;

	NICE::Matrix M( classNames.getMaxClassno() + 1, classNames.getMaxClassno() + 1 );
	M.set( 0 );
	for (int s = 0; s < ( int )M_vec.size(); s++ )
	{
		NICE::Matrix M_tmp = M_vec[s];
		for ( int r = 0; r < ( int )M_tmp.rows(); r++ )
		{
			for ( int c = 0; c < ( int )M_tmp.cols(); c++ )
			{
				if ( r == c )
					overall += M_tmp( r, c );

				sumall += M_tmp( r, c );
				M( r, c ) += M_tmp( 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 );
      double lsum = 0.0;
      for(int r2 = 0; r2 < ( int )M.rows(); r2++)
      {
        lsum += M(r,r2);
      }
      if(lsum != 0.0)
      {
        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();
  }

  fprintf( stderr, "overall: %f\n", overall );

  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;
}