#include "vislearning/features/localfeatures/LFColorWeijer.h"

#include <fstream>
#include <iostream>
#include "vislearning/baselib/ColorSpace.h"

using namespace OBJREC;

using namespace std;

using namespace NICE;

//! color representation for visualization
const int colors[11][3] =
{
  {0,     0,   0}, // black
  {0,     0, 255}, // blue
  {165,  42,  42}, // brown
  {190, 190, 190}, // grey
  {  0, 255,   0}, // green
  {255, 165,   0}, // orange
  {255, 192, 203}, // pink
  {160,  32, 240}, // purple
  {255,   0,   0}, // red
  {255, 255, 255}, // white
  {255, 255,   0}, // yellow
};


LFColorWeijer::LFColorWeijer( const Config *c )
{
  conf = c;

  bin[0] = conf->gI( "LFColorWeijer", "binL", 10 );
  bin[1] = conf->gI( "LFColorWeijer", "bina", 20 );
  bin[2] = conf->gI( "LFColorWeijer", "binb", 20 );

  maxv[0] =  100.0;
  maxv[1] =   80.0;
  maxv[2] =   50.0;

  minv[0] =    0.0;
  minv[1] = -105.0;
  minv[2] = -200.0;

  tfile = conf->gS( "LFColorWeijer", "table", "/home/dbv/bilder/colorWeijer/color.txt" );

  for ( int i = 0; i < 3; i++ )
  {
    interval[i] = ( maxv[i] - minv[i] ) / ( double )bin[i];
  }

  ifstream test( tfile.c_str() );

  if ( test )
  {
    restore();
  }
  else
  {
    train();
  }
}

LFColorWeijer::~LFColorWeijer()
{
  for ( uint i = 0; i < hist.size(); i++ )
  {
    for ( uint j = 0; j < hist[i].size(); j++ )
    {
      hist[i][j].clear();
    }
    hist[i].clear();
  }
  hist.clear();
}

int LFColorWeijer::getDescSize() const
{
  return LASTCOLOR;
}

void LFColorWeijer::store()
{
  ofstream fout( tfile.c_str(), ios_base::app );

  fout << hist.size() << " " << hist[0].size() << " " << hist[0][0].size() << " " << hist[0][0][0].size() << endl;

  for ( uint i = 0; i < hist.size(); i++ )
  {
    for ( uint i0 = 0; i0 < hist[i].size(); i0++ )
    {
      for ( uint i1 = 0; i1 < hist[i][i0].size(); i1++ )
      {
        for ( uint i2 = 0; i2 < hist[i][i0][i1].size(); i2++ )
        {
          fout << hist[i][i0][i1][i2] << " ";
        }
      }
    }
  }
}

void LFColorWeijer::smooth()
{
  int size0 = ( int )hist.size();
  int size1 = ( int )hist[0].size();
  int size2 = ( int )hist[0][0].size();
  int size3 = ( int )hist[0][0][0].size();
  for ( int i0 = 0; i0 < size1; i0++ )
  {
    for ( int i1 = 0; i1 < size2; i1++ )
    {
      for ( int i2 = 0; i2 < size3; i2++ )
      {
        double maxval = 0.0;
        for ( int i = 0; i < size0; i++ )
        {
          maxval = std::max( maxval, hist[i][i0][i1][i2] );
        }
        if ( maxval == 0.0 )
        {
          for ( int i = 0; i < size0; i++ )
          {
            int anz = 0;
            for ( int a0 = std::max( i0 - 1, 0 ); a0 <= std::min( i0 + 1, size1 - 1 ); a0++ )
            {
              for ( int a1 = std::max( i1 - 1, 0 ); a1 <= std::min( i1 + 1, size2 - 1 ); a1++ )
              {
                for ( int a2 = std::max( i2 - 1, 0 ); a2 <= std::min( i2 + 1, size3 - 1 ); a2++ )
                {
                  anz++;
                  hist[i][i0][i1][i2] += hist[i][a0][a1][a2];
                }
              }
            }
            hist[i][i0][i1][i2] /= anz;
          }
        }
      }
    }
  }
}

void LFColorWeijer::restore()
{
  int size0, size1, size2, size3;
  ifstream fin( tfile.c_str() );
  fin >> size0;
  fin >> size1;
  fin >> size2;
  fin >> size3;
  hist.clear();

  for ( int i = 0; i < size0; i++ )
  {
    vector<vector<vector<double> > > v2;

    for ( int i0 = 0; i0 < size1; i0++ )
    {
      vector<vector<double> > v1;

      for ( int i1 = 0; i1 < size2; i1++ )
      {
        vector<double> v0;

        for ( int i2 = 0; i2 < size3; i2++ )
        {
          double val;
          fin >> val;
          v0.push_back( val );
        }

        v1.push_back( v0 );
      }

      v2.push_back( v1 );
    }

    hist.push_back( v2 );
  }
}

int LFColorWeijer::getDescriptors( const NICE::Image & img, VVector & positions, VVector & features ) const
{
  cerr << "this are COLOR Features, they won't work on gray value images" << endl;
  exit( -1 );
}

int LFColorWeijer::getDescriptors( const NICE::ColorImage & img, VVector & positions, VVector & features ) const
{
  // in Lab umwandeln
  for ( int i = 0; i < ( int )positions.size(); i++ )
  {
    vector<double> vals;
    vector<int> b;
    int x = positions[i][0];
    int y = positions[i][1];

    double R, G, B, X, Y, Z;
    vector<double> lab( 3, 0.0 );

    R = ( double )img.getPixel( x, y, 0 ) / 255.0;
    G = ( double )img.getPixel( x, y, 1 ) / 255.0;
    B = ( double )img.getPixel( x, y, 2 ) / 255.0;

    ColorConversion::ccRGBtoXYZ( R, G, B, &X, &Y, &Z, 0 );
    ColorConversion::ccXYZtoCIE_Lab( X, Y, Z, &lab[0], &lab[1], &lab[2], 0 );

    for ( int i = 0; i < 3; i++ )
    {
      int val = ( int )(( lab[i] - minv[i] ) / interval[i] );
      val = std::min( val, bin[i] - 1 );
      val = std::max( val, 0 );
      b.push_back( val );
    }

    Vector feat( hist.size() );

    for ( uint i = 0; i < hist.size(); i++ )
    {
      feat[i] = hist[i][b[0]][b[1]][b[2]];
    }
    features.push_back( feat );
  }

  return 1;
}

void LFColorWeijer::visualizeFeatures( NICE::Image & mark, const VVector & positions, size_t color ) const
{

}

void LFColorWeijer::add( vector<vector<vector<double> > > &dest, vector<vector<vector<double> > > &src )
{
  for ( uint i0 = 0; i0 < src.size(); i0++ )
  {
    for ( uint i1 = 0; i1 < src[i0].size(); i1++ )
    {
      for ( uint i2 = 0; i2 < src[i0][i1].size(); i2++ )
      {
        dest[i0][i1][i2] += src[i0][i1][i2];
      }
    }
  }
}

int LFColorWeijer::findColor( string &fn )
{
  if ( fn.find( "black" ) != string::npos )
    return BLACK;
  if ( fn.find( "blue" ) != string::npos )
    return BLUE;
  if ( fn.find( "brown" ) != string::npos )
    return BROWN;
  if ( fn.find( "grey" ) != string::npos )
    return GREY;
  if ( fn.find( "green" ) != string::npos )
    return GREEN;
  if ( fn.find( "orange" ) != string::npos )
    return ORANGE;
  if ( fn.find( "pink" ) != string::npos )
    return PINK;
  if ( fn.find( "purple" ) != string::npos )
    return PURPLE;
  if ( fn.find( "red" ) != string::npos )
    return RED;
  if ( fn.find( "white" ) != string::npos )
    return WHITE;
  if ( fn.find( "yellow" ) != string::npos )
    return YELLOW;

  return -1;
}

vector<vector<vector<double > > > LFColorWeijer::createTable()
{
  vector<vector<vector<double> > > h;
  for ( int i0 = 0; i0 < bin[0]; i0++ )
  {
    vector<vector< double > > vec;
    for ( int i1 = 0; i1 < bin[1]; i1++ )
    {
      vector<double> v;
      for ( int i2 = 0; i2 < bin[2]; i2++ )
      {
        v.push_back( 0.0 );
      }
      vec.push_back( v );
    }
    h.push_back( vec );
  }
  return h;
}

void LFColorWeijer::normalize( vector<vector<vector<double> > > &tab )
{
  double sum = 0.0;

  for ( uint i0 = 0; i0 < tab.size(); i0++ )
  {
    for ( uint i1 = 0; i1 < tab[i0].size(); i1++ )
    {
      for ( uint i2 = 0; i2 < tab[i0][i1].size(); i2++ )
      {
        sum += tab[i0][i1][i2];
      }
    }
  }

  for ( uint i0 = 0; i0 < tab.size(); i0++ )
  {
    for ( uint i1 = 0; i1 < tab[i0].size(); i1++ )
    {
      for ( uint i2 = 0; i2 < tab[i0][i1].size(); i2++ )
      {
        tab[i0][i1][i2] /= sum;
      }
    }
  }

  return;
}

void LFColorWeijer::createHist( const ColorImage &cimg, vector<vector<vector<double> > > &hist, Image &mask )
{
  // in Lab umwandeln
  NICE::MultiChannelImageT<double> genimg, imglab;

  ColorSpace::ColorImagetoMultiChannelImage( cimg, genimg );
  ColorSpace::convert( imglab, genimg, ColorSpace::COLORSPACE_LAB, ColorSpace::COLORSPACE_RGB );

  for ( int y = 0; y < cimg.height(); y++ )
  {
    for ( int x = 0; x < cimg.width(); x++ )
    {
      if ( mask.getPixel( x, y ) == 0 )
        continue;
      vector<int> b;
      for ( int i = 0; i < 3; i++ )
      {
        int val = ( int )(( imglab.get( x, y, i ) - minv[i] ) / interval[i] );
        val = std::min( val, bin[i] - 1 );
        b.push_back( val );
      }
      hist[b[0]][b[1]][b[2]]++;
    }
  }
}

void LFColorWeijer::train()
{
  cout << "train Starts" << endl;
  for ( int i = 0; i < LASTCOLOR; i++ )
  {
    vector<vector<vector<double> > > h = createTable();
    hist.push_back( h );
  }

  string dir = conf->gS( "LFColorWeijer", "table", "/home/dbv/bilder/colorWeijer/ebay/" );
  string images = conf->gS( "LFColorWeijer", "table", "test_images.txt" );
  string mask = conf->gS( "LFColorWeijer", "table", "mask_images.txt" );

  string imagesfn;
  string maskfn;

  ifstream finimg(( dir + images ).c_str() );
  ifstream finmask(( dir + mask ).c_str() );
  cout << dir + images << endl;
  cout << dir + mask << endl;
  // lese bilder und masken ein
  while ( finimg >> imagesfn && finmask >> maskfn )
  {
    Image mimg( dir + maskfn );
    cout << dir + maskfn << endl;
    ColorImage cimg( dir + imagesfn );

    int col = findColor( imagesfn );
    vector<vector<vector<double> > > tab = createTable();

    createHist( cimg, tab, mimg ); // erzeuge Lab Histogramm des Bildes

    normalize( tab );

    add( hist[col], tab );
  }
  finimg.close();
  finmask.close();

  // normalisiere alle lookuptables
  for ( uint i = 0; i < hist.size(); i++ )
  {
    normalize( hist[i] );
  }

  smooth();
  store();
}

void LFColorWeijer::visualizeFeatures( NICE::ColorImage & out, const VVector & features, const VVector & position ) const
{
  for ( int i = 0; i < ( int )position.size(); i++ )
  {
    int maxpos = 0;
    double maxval = 0.0;
    for ( int j = 0; j < ( int )features[i].size(); j++ )
    {
      if ( maxval < features[i][j] )
      {
        maxval = features[i][j];
        maxpos = j;
      }
    }

    out.setPixel( position[i][0], position[i][1], colors[maxpos][0], colors[maxpos][1], colors[maxpos][2] );
  }
}

void LFColorWeijer::visualizeFeatures( const NICE::ColorImage & cimg ) const
{
  ColorImage out;
  visualizeFeatures( cimg, out );
}

void LFColorWeijer::visualizeFeatures( const NICE::ColorImage & cimg, NICE::ColorImage &out ) const
{
  VVector pos, feats;
  for ( int y = 0; y < cimg.height(); y++ )
  {
    for ( int x = 0; x < cimg.width(); x++ )
    {
      Vector vec( 2 );
      vec[0] = x;
      vec[1] = y;
      pos.push_back( vec );
    }
  }

  getDescriptors( cimg, pos, feats );
  
  //heatmap for a special class
  /*ofstream fout("out.txt");
  
  int counter = 0;
  for ( int y = 0; y < cimg.height(); y++ )
  {
    for ( int x = 0; x < cimg.width(); x++, counter++ )
    {
      fout << feats[counter][8] << " ";
    }
  }
  cout << "counter: " << counter << " feats.size(): " << feats.size() << endl;
  
  fout.close();*/

  out.resize( cimg.width(), cimg.height() );
  out.set( 0, 0, 0 );
  visualizeFeatures( out, feats, pos );
  ColorImage combinedout( cimg.width()*2, cimg.height() );

  int width = ( int )cimg.width();

  for ( int y = 0; y < ( int )cimg.height(); y++ )
  {
    for ( int x = 0; x < width; x++ )
    {
      combinedout.setPixel( x, y, cimg.getPixel( x, y, 0 ), cimg.getPixel( x, y, 1 ), cimg.getPixel( x, y, 2 ) );
      combinedout.setPixel( x + width, y, out.getPixel( x, y, 0 ), out.getPixel( x, y, 1 ), out.getPixel( x, y, 2 ) );
    }
  }

  showImage( combinedout, "result" );
}

void LFColorWeijer::getFeats( const ColorImage &img, MultiChannelImageT<double> &feats )
{
  int width = ( int )img.width();
  int height = ( int )img.height();
  feats.reInit( width, height, hist.size());

  NICE::MultiChannelImageT<double> genimg, imglab;

  ColorSpace::ColorImagetoMultiChannelImage( img, genimg );
  ColorSpace::convert( imglab, genimg, ColorSpace::COLORSPACE_LAB, ColorSpace::COLORSPACE_RGB );

  for ( int y = 0; y < height; y++ )
  {
    for ( int x = 0; x < width; x++ )
    {
      for ( uint i = 0; i < hist.size(); i++ )
      {
        vector<double> b( 3, 0.0 );
        for ( int j = 0; j < 3; j++ )
        {
          int val = ( int )(( imglab.get( x, y, j ) - minv[j] ) / interval[j] );
          val = std::min( val, bin[j] - 1 );
          val = std::max( val, 0 );
          b[j] = val;
        }
        feats.set( x, y, hist[i][b[0]][b[1]][b[2]], i );
      }
    }
  }

  return;
}