#include "RelativeLocationPrior.h"

#include "core/image/FilterT.h"

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

RelativeLocationPrior::RelativeLocationPrior()
{
  conf = new Config();
  mapsize = 200;
}

RelativeLocationPrior::RelativeLocationPrior ( const Config *_conf ) : conf ( _conf )
{
}

void RelativeLocationPrior::setClassNo ( int _classno )
{
  classno = _classno;
  Init();
}

void RelativeLocationPrior::Init()
{
  std::string section = "PostProcessRLP";
  mapsize = conf->gI ( section, "mapsize", 200 );

  featdim = classno * 3;

  //Priorsmaps erzeugen
  for ( int i = 0; i < classno; i++ )
  {
    NICE::MultiChannelImageT<double> *tmp  = new NICE::MultiChannelImageT<double> ( mapsize, mapsize, classno);
    tmp->setAll ( 0.0 );
    priormaps.push_back ( tmp );
  }
}

RelativeLocationPrior::~RelativeLocationPrior()
{
  for ( int i = 0; i < classno; i++ )
  {
    delete priormaps[i];
  }
}

void RelativeLocationPrior::trainPriorsMaps ( Examples &regions, int xsize, int ysize )
{
  for ( int j = 0; j < ( int ) regions.size(); j++ )
  {
    for ( int i = 0; i < ( int ) regions.size(); i++ )
    {
      if ( i == j )
        continue;

      int x = regions[i].second.x - regions[j].second.x;
      int y = regions[i].second.y - regions[j].second.y;

      convertCoords ( x, xsize );
      convertCoords ( y, ysize );

      priormaps[regions[i].first]->set ( x, y, priormaps[regions[i].first]->get ( x, y, regions[j].first ) + 1.0/*regions[j].second.weight*/, regions[j].first );
    }
  }
}

void RelativeLocationPrior::finishPriorsMaps ( ClassNames &cn )
{
  // Priormaps normalisieren
  double alpha = 5;
  for ( int i = 0; i < classno; i++ )
  {
    for ( int j = 0; j < classno; j++ )
    {
      double val = 0.0;

      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {
          val = std::max ( val, priormaps[i]->get ( x, y, j ) );
        }
      }
      if ( val != 0.0 )
      {
        for ( int x = 0; x < mapsize; x++ )
        {
          for ( int y = 0; y < mapsize; y++ )
          {
            double old = priormaps[i]->get ( x, y, j );

#undef DIRICHLET
#ifdef DIRICHLET
            old = ( old + alpha ) / ( val + classno * alpha );
#else
            old /= val;
#endif
            priormaps[i]->set ( x, y, old, j );
          }
        }
      }
    }
  }

  double sigma = 0.1 * ( double ) mapsize; // 10 percent of the maps height/width

  // Smoothing the maps
  for ( int j = 0; j < classno; j++ )
  {
    for ( int i = 0; i < classno; i++ )
    {
      NICE::FloatImage tmp ( mapsize, mapsize );
      tmp.set ( 0.0 );
      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {
          tmp.setPixelQuick ( x, y, priormaps[j]->get ( x, y, i ) );
        }
      }

      NICE::FloatImage out;
      NICE::FilterT<float,float,float> filter;
      filter.filterGaussSigmaApproximate( tmp, sigma, &out );

      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {
          priormaps[j]->set ( x, y, out.getPixel ( x, y ), i );
        }
      }
    }
  }

  // Sum of all pixels over all classes at a certain position equals 1
  for ( int i = 0; i < classno; i++ )
  {
    for ( int x = 0; x < mapsize; x++ )
    {
      for ( int y = 0; y < mapsize; y++ )
      {
        double val = 0.0;
        for ( int j = 0; j < classno; j++ )
        {
          val += priormaps[i]->get ( x, y, j );
        }
        if ( val != 0.0 )
        {
          for ( int j = 0; j < classno; j++ )
          {
            double old = priormaps[i]->get ( x, y, j );
            old /= val;
            priormaps[i]->set ( x, y, old, j );
          }
        }
      }
    }
  }

#undef VISDEBUG
#ifdef VISDEBUG
#ifndef NOVISUAL
  NICE::ColorImage rgbim ( ( classno - 1 ) * ( mapsize + 10 ), ( classno - 1 ) * ( mapsize + 10 ) );

  double maxval = -numeric_limits<double>::max();
  double minval = numeric_limits<double>::max();

  for ( int j = 0; j < classno; j++ )
  {
    if ( j == 6 ) continue;
    for ( int i = 0; i < classno; i++ )
    {
      if ( i == 6 ) continue;
      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {
          double val = priormaps[j]->get ( x, y, i );
          maxval = std::max ( val, maxval );
          minval = std::min ( val, minval );
        }
      }
    }
  }

  int jcounter = 0;
  for ( int j = 0; j < classno; j++ )
  {
    if ( j == 6 ) continue;
    int icounter = 0;
    for ( int i = 0; i < classno; i++ )
    {
      if ( i == 6 ) continue;

      NICE::FloatImage tmp ( mapsize, mapsize );
      tmp.set ( 0.0 );

      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {
          tmp.setPixel ( x, y, priormaps[j]->get ( x, y, i ) );
        }
      }

      tmp.setPixel ( 0, 0, maxval );
      tmp.setPixel ( 0, 1, minval );
      cout << "i: " << cn.text ( i ) << endl;
      NICE::ColorImage imgrgb2 ( mapsize, mapsize );
      ICETools::convertToRGB ( tmp, imgrgb2 );

      imgrgb2.setPixel ( 0, 0, 2, imgrgb2.getPixel ( 1, 0, 2 ) );
      imgrgb2.setPixel ( 0, 1, 2, imgrgb2.getPixel ( 1, 1, 2 ) );
      imgrgb2.setPixel ( 0, 0, 0, imgrgb2.getPixel ( 1, 0, 0 ) );
      imgrgb2.setPixel ( 0, 1, 0, imgrgb2.getPixel ( 1, 1, 0 ) );
      imgrgb2.setPixel ( 0, 0, 1, imgrgb2.getPixel ( 1, 0, 1 ) );
      imgrgb2.setPixel ( 0, 1, 1, imgrgb2.getPixel ( 1, 1, 1 ) );

      for ( int y = 0; y < mapsize; y++ )
      {
        for ( int x = 0; x < mapsize; x++ )
        {
          rgbim.setPixel ( x + jcounter* ( mapsize + 10 ), y + icounter* ( mapsize + 10 ), 2, imgrgb2.getPixel ( x, y, 2 ) );
          rgbim.setPixel ( x + jcounter* ( mapsize + 10 ), y + icounter* ( mapsize + 10 ), 0, imgrgb2.getPixel ( x, y, 0 ) );
          rgbim.setPixel ( x + jcounter* ( mapsize + 10 ), y + icounter* ( mapsize + 10 ), 1, imgrgb2.getPixel ( x, y, 1 ) );
        }
      }
      icounter++;
    }
    jcounter++;
  }
  rgbim.write ( "tmp.ppm" );
#endif
#endif
}

void RelativeLocationPrior::trainClassifier ( Examples &regions, NICE::MultiChannelImageT<double> & probabilities )
{
  // Creating a feature vector for all regions and adding it to the training set
  getFeature ( regions, probabilities );

  for ( int i = 0; i < ( int ) regions.size(); i++ )
  {
    trainingsdata.push_back ( pair<int, Example> ( regions[i].first, regions[i].second ) );
    regions[i].second.svec = NULL;
  }
}

void RelativeLocationPrior::finishClassifier()
{
  //////////////////////
  // Train Classifier //
  //////////////////////
  FeaturePool fp;
  Feature *f = new SparseVectorFeature ( featdim );
  f->explode ( fp );
  delete f;

  //feature size
  int s = 3;

  classifiers.resize ( classno );
  for ( int i = 0; i < classno; i++ )
  {
    classifiers[i] = SLR ( conf, "ClassifierSMLR" );
    Examples ex2;
    int countex = 0;
    for ( int j = 0; j < ( int ) trainingsdata.size(); j++ )
    {
      Example e;
      int z = 0;
      e.svec = new SparseVector ( s + 1 );
      for ( int k = i * s; k < i*s + s; k++, z++ )
      {
        double val = trainingsdata[j].second.svec->get ( k );
        if ( val != 0.0 )
          ( *e.svec ) [z] = val;
      }
      ( *e.svec ) [s] = 1.0;

      ex2.push_back ( pair<int, Example> ( trainingsdata[j].first, e ) );

      if ( trainingsdata[j].first == i )
        countex++;
    }

    if ( ex2.size() <= 2 || countex < 1 )
      continue;

    classifiers[i].train ( fp, ex2, i );

    for ( int j = 0; j < ( int ) ex2.size(); j++ )
    {
      delete ex2[j].second.svec;
      ex2[j].second.svec = NULL;
    }
  }

  trainingsdata.clear();
}

void RelativeLocationPrior::postprocess ( Examples &regions, NICE::MultiChannelImageT<double> & probabilities )
{
  getFeature ( regions, probabilities );

  int s = 3;

  for ( int i = 0; i < ( int ) regions.size(); i++ )
  {
    FullVector overall_distribution ( classno + 1 );
    overall_distribution[classno] = 0.0;

    double maxp = -numeric_limits<double>::max();
    int bestclass = 0;

    double sum  = 0.0;

    for ( int c = 0; c < classno; c++ )
    {
      Example e;
      int z = 0;
      e.svec = new SparseVector ( s + 1 );
      for ( int k = c * s; k < c*s + s; k++, z++ )
      {
        double val = regions[i].second.svec->get ( k );
        if ( val != 0.0 )
          ( *e.svec ) [z] = val;
      }
      ( *e.svec ) [s] = 1.0;

      overall_distribution[c] = classifiers[c].classify ( e );

      sum += overall_distribution[c];

      if ( maxp < overall_distribution[c] )
      {
        bestclass = c;
        maxp = overall_distribution[c];
      }
      delete e.svec;
      e.svec = NULL;
    }

    for ( int c = 0; c < classno; c++ )
    {
      overall_distribution[c] /= sum;
    }

    ClassificationResult r = ClassificationResult ( bestclass, overall_distribution );

    if ( bestclass < 0 )
    {
      regions[i].second.svec->store ( cout );
      cout << endl;
      cout << "fehler: besclass=" << bestclass << endl;
      for ( int j = 0; j < ( int ) probabilities.channels(); j++ )
      {
        cout << "j: " << j << " score: " << r.scores[j] << endl;
      }
    }
    regions[i].first = bestclass;
  }
}

void RelativeLocationPrior::convertCoords ( int &x, int xsize )
{
  x = ( int ) round ( ( double ( x ) + ( double ) xsize ) / ( 2.0 * ( double ) xsize ) * ( ( double ) mapsize - 1.0 ) );

  x = std::min ( x, mapsize - 1 );
  x = std::max ( x, 0 );
}

void RelativeLocationPrior::getFeature ( Examples &regions, NICE::MultiChannelImageT<double> & probabilities )
{

  int xsize, ysize;
  xsize = probabilities.width();
  ysize = probabilities.height();

  // get best classes
  vector<int> bestclasses ( regions.size(), -1 );
  for ( int r = 0; r < ( int ) regions.size(); r++ )
  {
    double maxval = -numeric_limits<double>::max();
    for ( int c = 0; c < ( int ) probabilities.channels(); c++ )
    {
      double val = probabilities.get ( regions[r].second.x, regions[r].second.y, c );
      if ( maxval < val )
      {
        bestclasses[r] = c;
        maxval = val;
      }
    }
  }

  vector<double> alpha;
  for ( int r = 0; r < ( int ) regions.size(); r++ )
  {
    double tmpalpha = probabilities.get ( regions[r].second.x, regions[r].second.y, bestclasses[r] ) * regions[r].second.weight;

    alpha.push_back ( tmpalpha );
  }

  // create f_relloc
  vector<vector<double> > vother;
  vector<vector<double> > vself;
  for ( int i = 0; i < ( int ) regions.size(); i++ )
  {
    vector<double> v, w;
    vother.push_back ( v );
    vself.push_back ( w );
    for ( int c = 0; c < classno; c++ )
    {
      double tmp_vother = 0.0;
      double tmp_self = 0.0;

      for ( int j = 0; j < ( int ) regions.size(); j++ )
      {
        if ( j == i )
          continue;

        int x = regions[i].second.x - regions[j].second.x;
        int y = regions[i].second.y - regions[j].second.y;

        convertCoords ( x, xsize );
        convertCoords ( y, ysize );

        double val = priormaps[c]->get ( x, y, bestclasses[j] ) * alpha[j]; ;

        if ( bestclasses[j] == bestclasses[i] ) //parts of the object
        {
          tmp_self += val;
        }
        else//Kontextinformationen
        {
          tmp_vother += val;
        }
      }

      if ( fabs ( tmp_self ) < 10e-7 )
        tmp_self = 10e-7;
      if ( fabs ( tmp_vother ) < 10e-7 )
        tmp_vother = 10e-7;

      vother[i].push_back ( tmp_vother );
      vself[i].push_back ( tmp_self );
    }
  }

  for ( int r = 0; r < ( int ) regions.size(); r++ )
  {
    if ( regions[r].second.svec != NULL )
    {
      delete regions[r].second.svec;
      regions[r].second.svec = NULL;
    }
    if ( regions[r].second.vec != NULL )
    {
      delete regions[r].second.vec;
      regions[r].second.vec = NULL;
    }

    regions[r].second.svec = new SparseVector ( classno*3 );

    int counter = 0;

    for ( int i = 0; i < classno; i++ )
    {
      //appearence feature (old probability for each class)
      double fapp = log ( probabilities.get ( regions[r].second.x, regions[r].second.y, i ) );

      if ( fabs ( fapp ) > 10e-7 )
        ( * ( regions[r].second.svec ) ) [counter] = fapp;
      counter++;

      double val = log ( vother[r][i] );

      if ( fabs ( val ) > 10e-7 )
        ( * ( regions[r].second.svec ) ) [counter] = val;
      counter++;

      val = log ( vself[r][i] );

      if ( fabs ( val ) > 10e-7 )
        ( * ( regions[r].second.svec ) ) [counter] = val;
      counter++;
    }
  }
}

void RelativeLocationPrior::restore ( istream & is, int format )
{
  is >> classno;
  is >> mapsize;
  is >> featdim;

  // Create prior maps
  for ( int i = 0; i < classno; i++ )
  {
    NICE::MultiChannelImageT<double> *tmp  = new NICE::MultiChannelImageT<double> ( mapsize, mapsize, classno);
    tmp->setAll ( 0.0 );
    priormaps.push_back ( tmp );
  }

  double val;
  for ( int i = 0; i < classno; i++ )
  {
    for ( int j = 0; j < classno; j++ )
    {
      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {

          is >> val;
          priormaps[i]->set ( x, y, val, j );
        }
      }
    }
  }

  classifiers.resize ( classno );
  for ( int i = 0; i < classno; i++ )
  {
    classifiers[i] = SLR();
    classifiers[i].restore ( is, format );
  }
}

void RelativeLocationPrior::store ( ostream & os, int format ) const
{
  os << classno << " ";
  os << mapsize << " ";
  os << featdim << endl;
  for ( int i = 0; i < classno; i++ )
  {
    for ( int j = 0; j < classno; j++ )
    {
      for ( int x = 0; x < mapsize; x++ )
      {
        for ( int y = 0; y < mapsize; y++ )
        {
          os << priormaps[i]->get ( x, y, j ) << " ";
        }
      }
    }
  }

  for ( int i = 0; i < classno; i++ )
  {
    classifiers[i].store ( os, format );
  }
}

void RelativeLocationPrior::clear ()
{

}