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

/**
* @file testSemanticSegmentation.cpp
* @brief test semantic segmentation routines for 3d images and 2d images
* @author Erik Rodner, Björn Fröhlich, Sven Sickert
* @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 <core/image/MultiChannelImage3DT.h>
#include <semseg3d/semseg/SemSegContextTree3D.h>

#include <core/basics/ResourceStatistics.h>
#include <core/image/Morph.h>

#include <fstream>
#include <vector>

#undef DEBUG

using namespace OBJREC;

using namespace NICE;

using namespace std;

void segmentToOverlay ( const NICE::Image *orig, const NICE::ColorImage & segment,
                        NICE::ColorImage & result )
{
  int xsize = orig->width();
  int ysize = orig->height();
  
  result.resize( xsize, ysize );
  vector< NICE::MatrixT<double> > channelMat;
  
  double alpha = .5;
  
  for (int c = 0; c < 3; c++)
  {
    NICE::MatrixT<double> chan ( xsize, ysize );
    channelMat.push_back( chan );
  }

  for (int y = 0; y < ysize; y++)
  {
    for (int x = 0; x < xsize; x++)
    {
      uchar val = orig->getPixelQuick(x,y);
      for (int c = 0; c < 3; c++)
        channelMat[c](x,y) = (double)val + alpha*(double)segment.getPixel( x, y, c );
    }
  }
  
  for (int c = 0; c < 3; c++)
  {
    channelMat[c] /= channelMat[c].Max();
    channelMat[c] *= 255;
  }
  
  for (int y = 0; y < ysize; y++)
  {
    for (int x = 0; x < xsize; x++)
    {
      for (int c = 0; c < 3; c++)
      {
        int val = channelMat[c](x,y);
        result.setPixel( x, y, c, (uchar)val);
      }
    }
  }
}

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;

  /*-------------I/O CONFIGURATION-------------*/
  bool postProcessing = conf.gB( "main", "post_process", false);
  bool run_3Dseg = conf.gB( "SSContextTree", "run_3dseg", false);
  bool show_result = conf.gB ( "debug", "show_results", false );
  bool write_results = conf.gB ( "debug", "write_results", false );
  string output_type = conf.gS ( "debug", "output_type", "ppm" );
  string output_postfix = conf.gS ( "debug", "output_postfix", "" );
  string resultdir = conf.gS ( "debug", "resultdir", "." );
  /*-------------------------------------------*/

#ifdef DEBUG
  cerr << "Writing Results to " << resultdir << endl;
#endif


  MultiDataset md ( &conf );

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

  // initialize semantic segmentation method
  SemanticSegmentation *semseg = NULL;
  semseg = new SemSegContextTree3D ( &conf, &md );
  
  // train semantic segmentation method
  cout << "\nTRAINING" << endl;
  cout << "########\n" << endl;
  semseg->train( &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();

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

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

  cout << "\nCLASSIFICATION" << endl;
  cout << "##############\n" << endl;
  for (LabeledSet::const_iterator it = testFiles->begin(); it != testFiles->end(); it++)
  {
    for (std::vector<ImageInfo *>::const_iterator jt = it->second.begin();
         jt != it->second.end(); jt++)
    {
      ImageInfo & info = *(*jt);
      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() );
#ifdef DEBUG
        cout << "testSemanticSegmentation3D: Generating Labeled NICE::Image (Ground-Truth)" << endl;
#endif
        l_gt->calcLabeledImage ( lm_gt, classNames.getBackgroundClass() );
      }
      segresult.addChannel ( lm );
      gt.addChannel ( lm_gt );

      int depthBoundary = 0;
      if ( run_3Dseg )
      {
        depthBoundary = zsizeVec[idx];
      }

      if ( depthCount < depthBoundary ) continue;

      NICE::MultiChannelImage3DT<double> probabilities;
      NICE::MultiChannelImage3DT<double> imgData;

      semseg->make3DImage ( filelist, imgData );
      semseg->classify ( imgData, segresult, probabilities, filelist );

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

        if ( show_result || write_results )
        {
          NICE::ColorImage orig ( filelist[z] );
          NICE::ColorImage rgb;
          NICE::ColorImage rgb_gt;
          NICE::ColorImage ov_rgb;
          NICE::ColorImage ov_rgb_gt;

          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 ) );
              if ( run_3Dseg )
                lm_gt.setPixel ( x, y, gt.get ( x, y, ( uint ) z ) );
            }
          }

          // confusion matrix
          NICE::Matrix M ( classNames.getMaxClassno() + 1, classNames.getMaxClassno() + 1 );
          M.set ( 0 );
          updateMatrix ( lm, lm_gt, M, forbidden_classes );
          M_vec.push_back ( M );

          classNames.labelToRGB ( lm, rgb );
          classNames.labelToRGB ( lm_gt, rgb_gt );

          if (postProcessing)
          {
            // median filter
            for (int r = 0; r < 3; r++)
            {
              NICE::Image postIm(rgb.width(), rgb.height());
              NICE::median(*(rgb.getChannel(r)), &postIm, 1);
              for (int y = 0; y < rgb.height(); y++)
                for (int x = 0; x < rgb.width(); x++)
                  rgb.setPixel(x,y,r, postIm.getPixelQuick(x,y));
            }
          }

          segmentToOverlay ( orig.getChannel(1), rgb, ov_rgb );
          segmentToOverlay ( orig.getChannel(1), rgb_gt, ov_rgb_gt );
          
          if ( write_results )
          {
            std::stringstream out;       
            if ( output_postfix.size() > 0 )
              out << resultdir << "/" << fname << output_postfix;
            else
              out << resultdir << "/" << fname;
#ifdef DEBUG
            cout << "Writing to file " << out.str() << "_*." << output_type << endl;
#endif
            orig.write ( out.str() + "_orig." + output_type );
            rgb.write ( out.str() + "_result." + output_type );
            rgb_gt.write ( out.str() + "_groundtruth." + output_type );
            ov_rgb.write ( out.str() + "_overlay_res." + output_type );
            ov_rgb_gt.write ( out.str() + "_overlay_gt." + output_type );
          }
        }
      }

      // prepare for new 3d image
      filelist.clear();
      segresult.reInit(0,0,0);
      gt.reInit(0,0,0);
      depthCount = 0;
      idx++;
//       pb.update ( testFiles->count() );
    }
  }
  
  segresult.freeData();
//   pb.hide();

  cout << "\nSTATISTICS" << endl;
  cout << "##########\n" << endl;
  long maxMemory;
  double userCPUTime, sysCPUTime;
  rs.getStatistics ( maxMemory, userCPUTime, sysCPUTime );
  cout << "Memory (max):    " << maxMemory << " KB" << endl;
  cout << "CPU Time (user): " << userCPUTime << " seconds" << endl;
  cout << "CPU Time (sys):  " << sysCPUTime << " seconds" << 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;
  }

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

  // print/save results of evaluation
  cout << "\nPERFORMANCE" << endl;
  cout << "###########\n" << endl;
  ofstream fout ( ( resultdir + "/res.txt" ).c_str(), ios::out );
  fout << "Overall Recognition Rate: " << overall << endl;
  fout << "Average Recognition Rate: " << avg_perf / ( classes_trained ) << endl;
  fout << "Lower Bound: " << 1.0  / classes_trained << endl;
  cout << "Overall Recogntion Rate: " << overall << endl;
  cout << "Average Recogntion Rate: " << avg_perf / ( classes_trained ) << endl;
  cout << "Lower Bound: " << 1.0  / classes_trained << endl;

  cout <<"\nClasses:" << 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;
      cout << classname.c_str() << ": " << M ( r, r ) << endl;
    }
  }
  fout.close();

  delete semseg;

  return 0;
}