#include <sstream>
#include <iostream>

#include "SemSegNovelty.h"

#include "core/image/FilterT.h"
#include "gp-hik-exp/GPHIKClassifierNICE.h"
#include "vislearning/baselib/ICETools.h"
#include "vislearning/baselib/Globals.h"
#include "vislearning/features/fpfeatures/SparseVectorFeature.h"
#include "core/basics/StringTools.h"
#include "core/basics/Timer.h"

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

SemSegNovelty::SemSegNovelty ( const Config *conf,
                               const MultiDataset *md )
    : SemanticSegmentation ( conf, & ( md->getClassNames ( "train" ) ) )
{
  this->conf = conf;

  string section = "SemSegNovelty";

  featExtract = new LFColorWeijer ( conf );

  save_cache = conf->gB ( "FPCPixel", "save_cache", true );
  read_cache = conf->gB ( "FPCPixel", "read_cache", false );
  uncertdir = conf->gS("debug", "uncertainty", "uncertainty");
  cache = conf->gS ( "cache", "root", "" );


  classifier = new GPHIKClassifierNICE ( conf, "ClassiferGPHIK" );;

  whs = conf->gI ( section, "window_size", 10 );
  featdist = conf->gI ( section, "grid", 10 );
  testWSize = conf->gI (section, "test_window_size", 10);

  cn = md->getClassNames ( "train" );

  if ( read_cache )
  {
    string classifierdst = "/classifier.data";
    fprintf ( stderr, "SemSegNovelty:: Reading classifier data from %s\n", ( cache + classifierdst ).c_str() );

    try
    {
      if ( classifier != NULL )
      {
        classifier->read ( cache + classifierdst );
      }

      fprintf ( stderr, "SemSegNovelty:: successfully read\n" );
    }
    catch ( char *str )
    {
      cerr << "error reading data: " << str << endl;
    }
  }
  else
  {
    train ( md );
  }
}

SemSegNovelty::~SemSegNovelty()
{
  // clean-up
  if ( classifier != NULL )
    delete classifier;
  if ( featExtract != NULL )
    delete featExtract;
}

void SemSegNovelty::train ( const MultiDataset *md )
{
  const LabeledSet train = * ( *md ) ["train"];
  const LabeledSet *trainp = &train;

  ////////////////////////
  // feature extraction //
  ////////////////////////

  std::string forbidden_classes_s = conf->gS ( "analysis", "donttrain", "" );
  if ( forbidden_classes_s == "" )
  {
    forbidden_classes_s = conf->gS ( "analysis", "forbidden_classes", "" );
  }
  cn.getSelection ( forbidden_classes_s, forbidden_classes );
  cerr << "forbidden: " << forbidden_classes_s << endl;

  ProgressBar pb ( "Local Feature Extraction" );
  pb.show();

  int imgnb = 0;

  Examples examples;
  examples.filename = "training";

  int featdim = -1;

  LOOP_ALL_S ( *trainp )
  {
    //EACH_S(classno, currentFile);
    EACH_INFO ( classno, info );

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

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

    Image labels ( xsize, ysize );
    labels.set ( 0 );
    locResult->calcLabeledImage ( labels, ( *classNames ).getBackgroundClass() );

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

    Globals::setCurrentImgFN ( currentFile );

    MultiChannelImageT<double> feats;

    // extract features
    featExtract->getFeats ( img, feats );
    featdim = feats.channels();
    feats.addChannel(featdim);

    for (int c = 0; c < featdim; c++)
    {
      ImageT<double> tmp = feats[c];
      ImageT<double> tmp2 = feats[c+featdim];

      NICE::FilterT<double, double, double>::gradientStrength (tmp, tmp2);
    }
    featdim += featdim;

    // compute integral images
    for ( int c = 0; c < featdim; c++ )
    {
      feats.calcIntegral ( c );
    }

    for ( int y = 0; y < ysize; y += featdist )
    {
      for ( int x = 0; x < xsize; x += featdist )
      {
        int classno = labels ( x, y );

        if ( forbidden_classes.find ( classno ) != forbidden_classes.end() )
          continue;


        Example example;
        example.vec = NULL;
        example.svec = new SparseVector ( featdim );
        for ( int f = 0; f < featdim; f++ )
        {
          double val = feats.getIntegralValue ( x - whs, y - whs, x + whs, y + whs, f );
          if ( val > 1e-10 )
            ( *example.svec ) [f] = val;
        }

        example.svec->normalize();

        example.position = imgnb;
        examples.push_back ( pair<int, Example> ( classno, example ) );
      }
    }

    delete ce;
    imgnb++;
    pb.update ( trainp->count() );
  }

  pb.hide();


  //////////////////////
  // train classifier //
  //////////////////////
  FeaturePool fp;

  Feature *f = new SparseVectorFeature ( featdim );

  f->explode ( fp );
  delete f;

  if ( classifier != NULL )
    classifier->train ( fp, examples );
  else
  {
    cerr << "no classifier selected?!" << endl;
    exit ( -1 );
  }

  fp.destroy();

  if ( save_cache )
  {
    if ( classifier != NULL )
      classifier->save ( cache + "/classifier.data" );
  }

  ////////////
  //clean up//
  ////////////
  for ( int i = 0; i < ( int ) examples.size(); i++ )
  {
    examples[i].second.clean();
  }
  examples.clear();

  cerr << "SemSeg training finished" << endl;
}

void SemSegNovelty::semanticseg ( CachedExample *ce, NICE::Image & segresult, NICE::MultiChannelImageT<double> & probabilities )
{
  Timer timer;
  timer.start();

  Examples examples;
  examples.filename = "testing";

  segresult.set ( 0 );

  int featdim = -1;

  std::string currentFile = Globals::getCurrentImgFN();


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

  probabilities.reInit( xsize, ysize, cn.getMaxClassno() + 1);
  probabilities.set ( 0.0 );

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

  MultiChannelImageT<double> feats;

  // extract features
  featExtract->getFeats ( img, feats );
  featdim = feats.channels();
  feats.addChannel(featdim);

  for (int c = 0; c < featdim; c++)
  {
    ImageT<double> tmp = feats[c];
    ImageT<double> tmp2 = feats[c+featdim];

    NICE::FilterT<double, double, double>::gradientStrength (tmp, tmp2);
  }
  featdim += featdim;

  // compute integral images
  for ( int c = 0; c < featdim; c++ )
  {
    feats.calcIntegral ( c );
  }

  FloatImage uncert ( xsize, ysize );
  uncert.set ( 0.0 );

  double maxunc = -numeric_limits<double>::max();
  timer.stop();
  cout << "first: " << timer.getLastAbsolute() << endl;
  timer.start();
#pragma omp parallel for
  for ( int y = 0; y < ysize; y += testWSize )
  {
    Example example;
    example.vec = NULL;
    example.svec = new SparseVector ( featdim );
    for ( int x = 0; x < xsize; x += testWSize)
    {
      for ( int f = 0; f < featdim; f++ )
      {
        double val = feats.getIntegralValue ( x - whs, y - whs, x + whs, y + whs, f );
        if ( val > 1e-10 )
          ( *example.svec ) [f] = val;
      }
      example.svec->normalize();

      ClassificationResult cr = classifier->classify ( example );

      int xs = std::max(0, x - testWSize/2);
      int xe = std::min(xsize - 1, x + testWSize/2);
      int ys = std::max(0, y - testWSize/2);
      int ye = std::min(ysize - 1, y + testWSize/2);
      for (int yl = ys; yl <= ye; yl++)
      {
        for (int xl = xs; xl <= xe; xl++)
        {
          for ( int j = 0 ; j < cr.scores.size(); j++ )
          {
            probabilities ( xl, yl, j ) = cr.scores[j];
          }
          segresult ( xl, yl ) = cr.classno;
          uncert ( xl, yl ) = cr.uncertainty;
        }
      }

      if (maxunc < cr.uncertainty)
        maxunc = cr.uncertainty;
      example.svec->clear();
    }
    delete example.svec;
    example.svec = NULL;
  }

  cout << "maxunertainty: " << maxunc << endl;

  timer.stop();
  cout << "second: " << timer.getLastAbsolute() << endl;
  timer.start();

  ColorImage imgrgb ( xsize, ysize );

  std::stringstream out;
  std::vector< std::string > list2;
  StringTools::split ( Globals::getCurrentImgFN (), '/', list2 );
  out << uncertdir << "/" << list2.back();

  uncert.writeRaw(out.str() + ".rawfloat");
  uncert(0, 0) = 0.0;
  uncert(0, 1) = 1.0;
  ICETools::convertToRGB ( uncert, imgrgb );
  imgrgb.write ( out.str() + "rough.png" );


  timer.stop();
  cout << "last: " << timer.getLastAbsolute() << endl;
}