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

  /*-------------I/O CONFIGURATION-------------*/
  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 );
  bool run_3dseg = conf.gB ( "debug", "run_3dseg", true );
  string output_type = conf.gS ( "debug", "output_type", "ppm" );
  string output_postfix = conf.gS ( "debug", "output_postfix", "" );
  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 );

    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 );
    cout << segresult.width() << " " << segresult.height() << " " << segresult.channels() << endl;
    semseg->semanticseg ( imgData, segresult, probabilities, filelist );

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

    // save to file
    for ( int z = 0; z < segresult.channels(); 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];
        sprintf ( filename, "%s/%s.%s.%s", resultdir.c_str(), fname.c_str(), output_postfix.c_str(), output_type.c_str() );
        pascal_lm.write ( filename );
      }

      if ( show_result || write_results )
      {
        NICE::ColorImage orig ( filelist[z] );
        NICE::ColorImage rgb;
        NICE::ColorImage 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 ) );
          }
        }

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

        if ( write_results )
        {
          char filename[1024];
          if ( output_postfix.size() > 0 )
            sprintf ( filename, "%s.%s.%s", fname.c_str(), output_postfix.c_str(), output_type.c_str() );
          else
            sprintf ( filename, "%s.%s", fname.c_str(), output_type.c_str() );
          
          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 < segresult.channels(); z++ )
    {
      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 ) );
        }
      }
      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();
    segresult.reInit(0,0,0);
    gt.reInit(0,0,0);
    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;
}