#include "FeatureLearningPrototypes.h"

//STL
#include <iostream>

//core
#include <core/image/FilterT.h>
#include <core/image/CircleT.h>
#include <core/image/Convert.h>
#include <core/vector/VectorT.h>

//vislearning
#include <vislearning/features/localfeatures/LFonHSG.h>
#include <vislearning/features/localfeatures/LFColorSande.h>
#include <vislearning/features/localfeatures/LFColorWeijer.h>
#include <vislearning/features/localfeatures/LFReadCache.h>
#include <vislearning/features/localfeatures/LFWriteCache.h>
// 
#include <vislearning/math/cluster/KMeans.h>
#include <vislearning/math/cluster/GMM.h>

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

  //**********************************************
  //
  //                 PROTECTED METHODS
  //
  //********************************************** 

void FeatureLearningPrototypes::setClusterAlgo( const std::string & _clusterAlgoString)
{
  //be careful with previously allocated memory
  if (this->clusterAlgo != NULL)
    delete clusterAlgo;
  
  if (_clusterAlgoString.compare("kmeans") == 0)
  {
    this->clusterAlgo = new OBJREC::KMeans(this->initialNumberOfClusters);
  }
  else if (_clusterAlgoString.compare("GMM") == 0) 
  {
    this->clusterAlgo = new OBJREC::GMM(this->conf, this->initialNumberOfClusters);
  }
  else
  {
    std::cerr << "Unknown cluster algorithm selected, use k-means instead" << std::endl;
    this->clusterAlgo = new OBJREC::KMeans(this->initialNumberOfClusters);
  }    
}

void FeatureLearningPrototypes::setFeatureExtractor( const bool & _setForTraining )
{  
  //be careful with previously allocated memory
  if (this->featureExtractor != NULL)
    delete featureExtractor;  
  
    //feature stuff
  // which OpponentSIFT implementation to use {NICE, VANDESANDE}
  std::string opSiftImpl;  
  opSiftImpl = this->conf->gS ( "Descriptor", "implementation", "VANDESANDE" );
  // read features?
  bool readfeat;
  readfeat = this->conf->gB ( "Descriptor", "read", true );
  // write features?
  bool writefeat;  
  writefeat = this->conf->gB ( "Descriptor", "write", true );   
  
  // Welche Opponentsift Implementierung soll genutzt werden ?
  LocalFeatureRepresentation *cSIFT = NULL;
  LocalFeatureRepresentation *writeFeats = NULL;
  LocalFeatureRepresentation *readFeats = NULL;
  this->featureExtractor = NULL;
  if ( opSiftImpl == "NICE" )
  {
    if ( _setForTraining )
      cSIFT = new OBJREC::LFonHSG ( this->conf, "HSGtrain" );
    else
      cSIFT = new OBJREC::LFonHSG ( this->conf, "HSGtest" );
  }
  else if ( opSiftImpl == "VANDESANDE" )
  {
    if ( _setForTraining )
      cSIFT = new OBJREC::LFColorSande ( this->conf, "LFColorSandeTrain" );
    else
      cSIFT = new OBJREC::LFColorSande ( this->conf, "LFColorSandeTest" );
  }
  else
  {
    fthrow ( Exception, "feattype: %s not yet supported" << opSiftImpl );
  }

  this->featureExtractor = cSIFT;
  
  if ( writefeat )
  {
    // write the features to a file, if there isn't any to read
    writeFeats = new LFWriteCache ( this->conf, cSIFT );
    this->featureExtractor = writeFeats;
  }

  if ( readfeat )
  {
    // read the features from a file
    if ( writefeat )
    {
      readFeats = new LFReadCache ( this->conf, writeFeats, -1 );
    }
    else
    {
      readFeats = new LFReadCache ( this->conf, cSIFT, -1 );
    }
    this->featureExtractor = readFeats; 
  }  
  
  //only set feature stuff to NULL, deletion of the underlying object is done in the destructor
  if ( cSIFT != NULL )
    cSIFT = NULL;
  if ( writeFeats != NULL )
    writeFeats = NULL;
  if ( readFeats != NULL )
    readFeats  = NULL ;   
}

void FeatureLearningPrototypes::extractFeaturesFromTrainingImages( const OBJREC::MultiDataset *_md, NICE::VVector & examplesTraining )
{
  examplesTraining.clear();
  
  int numberOfTrainImage ( 0 );
  
  const LabeledSet *trainFiles = (*_md)["train"]; 
  
  //run over all training images
  LOOP_ALL_S( *trainFiles )
  {
      EACH_INFO( classno, info );
      std::string filename = info.img();
            
      NICE::ColorImage img( filename );
      if ( b_showTrainingImages )
      {
        showImage( img, "Input" );    
      }
      
      //variables to store feature informatio
      NICE::VVector features;
      NICE::VVector cfeatures;
      NICE::VVector positions;

      //compute features
      Globals::setCurrentImgFN ( filename );
      if (featureExtractor == NULL)
        std::cerr << "feature Extractor is NULL" << std::endl;
      else
        featureExtractor->extractFeatures ( img, features, positions );
            
      //store feature information in larger data structure
      for ( NICE::VVector::iterator i = features.begin();
            i != features.end();
            i++)
      {              
        //normalization :)
        i->normalizeL1();

        examplesTraining.push_back(*i);
      }
      
      //don't waste memory
      features.clear();
      positions.clear();      
      numberOfTrainImage++;
  }//Loop over all training images    
}

void FeatureLearningPrototypes::train ( const OBJREC::MultiDataset *_md )
{  
  bool loadSuccess = this->loadInitialCodebook();
  
  if ( !loadSuccess )
  {
    //**********************************************
    //
    //     EXTRACT FEATURES FROM TRAINING IMAGES
    //
    //**********************************************  
    
    std::cerr << " EXTRACT FEATURES FROM TRAINING IMAGES" << std::endl;
    
    NICE::VVector examplesTraining;  
    this->extractFeaturesFromTrainingImages( _md, examplesTraining  );
    
    //**********************************************
    //
    //    CLUSTER FEATURES FROM TRAINING IMAGES
    //
    //    THIS GIVES US AN INITIAL CODEBOOK
    //
    //**********************************************  
    std::cerr << " CLUSTER FEATURES FROM TRAINING IMAGES" << std::endl;
    //go, go, go...  
    prototypes.clear();
    std::vector< double > weights;
    std::vector< int > assignment;
    clusterAlgo->cluster ( examplesTraining, prototypes, weights, assignment);
    weights.clear();
    assignment.clear();  
    
    //normalization
    for (NICE::VVector::iterator protoIt = prototypes.begin(); protoIt != prototypes.end(); protoIt++)
    {
      protoIt->normalizeL1();
    }
  }
  
  this->writeInitialCodebook();
}

bool FeatureLearningPrototypes::loadInitialCodebook ( )
{
  if ( b_loadInitialCodebook  )
  {
    std::cerr << " INITIAL CODEBOOK ALREADY COMPUTED - RE-USE IT" << std::endl;    
    std::cerr << " // WARNING - WE DO NOT VERIFY WHETHER THIS IS THE CORRECT CODEBOOK FOR THIS TRAINING SET!!!!" << std::endl;    
    
    prototypes.clear();
    
    try
    {
      prototypes.read(cacheInitialCodebook);
    }
    catch (...)
    {
      std::cerr << "Error while loading initial codebook" << std::endl;
      return false;
    }  
    return true;
  }
  else
    return false;
}

bool FeatureLearningPrototypes::writeInitialCodebook ( )
{      
  if (  b_saveInitialCodebook )
  {
    std::cerr << " SAVE INITIAL CODEBOOK " << std::endl;
    
    try 
    {
      prototypes.write( cacheInitialCodebook );
    }
    catch (...)
    {
      std::cerr << "Error while saving initial codebook" << std::endl;
      return false;
    }    
    return true;
  } 
  else
    return false;
}


void FeatureLearningPrototypes::evaluateCurrentCodebookForGivenFeatures(  const NICE::VVector & _features,
                                                                          const NICE::VVector & _positions,
                                                                          NICE::FloatImage & _noveltyImageGaussFiltered,
                                                                          NICE::FloatImage * _noveltyImage                                                                        )
{
  bool wasNoveltyImageGiven ( true );
  if ( _noveltyImage == NULL )
  {
    _noveltyImage = new FloatImage ( _noveltyImageGaussFiltered.width(), _noveltyImageGaussFiltered.height() );
    wasNoveltyImageGiven = false;
  }
  
  _noveltyImageGaussFiltered.set( 0.0 );
  _noveltyImage->set( 0.0 );
     
  
  NICE::VVector::const_iterator posIt = _positions.begin();
  for ( NICE::VVector::const_iterator i = _features.begin();
        i != _features.end();
        i++, posIt++)
  {              
    
    //loop over codebook representatives
    double minDist ( std::numeric_limits<double>::max() );
    for (NICE::VVector::const_iterator it =  prototypes.begin(); it != prototypes.end(); it++)
    {
      //compute distance
      double tmpDist ( this->distFunction->calculate(*i,*it) );
      if (tmpDist < minDist)
        minDist = tmpDist;
    }

    //take minimum distance and store in in a float image    
    (*_noveltyImage) ( (*posIt)[0], (*posIt)[1]  ) = minDist;
  } 

  
  //gauss-filtering for nicer visualization
  float sigma ( 3.0 );
  FilterT<float, float, float> filter;
  filter.filterGaussSigmaApproximate ( *_noveltyImage, sigma, & _noveltyImageGaussFiltered );
    
  if ( ! wasNoveltyImageGiven )
    delete _noveltyImage;
}

  //**********************************************
  //
  //                 PUBLIC METHODS
  //
  //********************************************** 


FeatureLearningPrototypes::FeatureLearningPrototypes ( const Config *_conf,
                               const MultiDataset *_md, const std::string & _section )
    : FeatureLearningGeneric ( _conf, _section )
{ 
  
  // define the initial number of clusters our codebook shall contain
  initialNumberOfClusters = conf->gI(section, "initialNumberOfClusters", 10);
  
  // define the clustering algorithm to be used
  std::string clusterAlgoString = conf->gS(section, "clusterAlgo", "kmeans");
  
  // define the distance function to be used
  std::string distFunctionString = conf->gS(section, "distFunction", "euclidian");    
  
  
  //**********************************************
  //
  //      SET UP VARIABLES AND METHODS
  //             - FEATURE TYPE
  //             - CLUSTERING ALGO
  //             - DISTANCE FUNCTION
  //             - ...
  //
  //**********************************************  
  
  std::cerr << " SET UP VARIABLES AND METHODS " << std::endl;

  //feature extraction for initial codebook
  this->featureExtractor = NULL;
  this->setFeatureExtractor( true /* set for training */ );  
  
  //clustering algorithm
  this->clusterAlgo = NULL;
  this->setClusterAlgo( clusterAlgoString );
  
  if (distFunctionString.compare("euclidian") == 0)
  {
    distFunction = new NICE::EuclidianDistance<double>();
  }
  else
  {
    std::cerr << "Unknown vector distance selected, use euclidian instead" << std::endl;
    distFunction = new NICE::EuclidianDistance<double>();
  }    
  
  //run the training to initially compute a codebook and stuff like that
  this->train( _md );  
    
  //so far, we have not seen any new image
  this->newImageCounter = 0;
  
  //TODO stupid
  this->maxValForVisualization = conf->gD( section, "stupidMaxValForVisualization", 0.005 ) ;
  
  
  //feature extraction for unseen images
  this->setFeatureExtractor( false /* set for training */ );    
}

FeatureLearningPrototypes::~FeatureLearningPrototypes()
{
  // clean-up
  if ( clusterAlgo != NULL )
    delete clusterAlgo;
  if ( distFunction != NULL )
    delete distFunction;
  if ( featureExtractor != NULL )
    delete featureExtractor; 
}

NICE::FloatImage FeatureLearningPrototypes::evaluateCurrentCodebookByDistance ( const std::string & _filename , const bool & beforeComputingNewFeatures )
{
  std::cerr << " VISUALIZATION -----    maxValForVisualization: " << maxValForVisualization << std::endl;
  
    NICE::ColorImage img( _filename );
    if ( b_showTrainingImages )
    {
      showImage( img, "Input" );    
    }
    
    int xsize ( img.width() );
    int ysize ( img.height() );
    
    //variables to store feature information
    NICE::VVector features;
    NICE::VVector cfeatures;
    NICE::VVector positions;

    //compute features
    Globals::setCurrentImgFN ( _filename );
    featureExtractor->extractFeatures ( img, features, positions );
    
    //normalization
    for ( NICE::VVector::iterator i = features.begin();
          i != features.end();
          i++)
    {              
      //normalization :)
      i->normalizeL1();
    }
    
    FloatImage noveltyImage ( xsize, ysize );
    FloatImage noveltyImageGaussFiltered ( xsize, ysize );
    
    this->evaluateCurrentCodebookForGivenFeatures( features, positions, noveltyImageGaussFiltered, &noveltyImage );
    
    double maxDist ( noveltyImage.max() );
    double maxFiltered ( noveltyImageGaussFiltered.max() );

    std::cerr << "maximum distance of Training images: " << maxDist << std::endl;  
    std::cerr << "maximum distance of Training images after filtering: " << maxFiltered << std::endl;      
    if ( beforeComputingNewFeatures )
      this->oldMaxDist = maxFiltered;

    
    //convert float to RGB
    NICE::ColorImage noveltyImageRGB ( xsize, ysize  );
//     ICETools::convertToRGB ( noveltyImageGaussFiltered, noveltyImageRGB );
    if ( beforeComputingNewFeatures )
    {
      imageToPseudoColorWithRangeSpecification( noveltyImageGaussFiltered, noveltyImageRGB, 0 /* min */, maxValForVisualization /* maxFiltered*/ /* max */ );
      std::cerr << "set max value to: " << noveltyImageGaussFiltered.max() << std::endl;
    }
    else
    {
      imageToPseudoColorWithRangeSpecification( noveltyImageGaussFiltered, noveltyImageRGB, 0 /* min */, maxValForVisualization /*this->oldMaxDist*/ /* max */ );
      std::cerr << "set max value to: " << this->oldMaxDist << std::endl;
    }
    
    
    
    if ( b_showResults )
      showImage(noveltyImageRGB, "Novelty Image");  
    else 
    {
      std::vector< std::string > list2;
      StringTools::split ( _filename, '/', list2 );      

      std::string destination ( s_resultdir + NICE::intToString(this->newImageCounter -1 ) + "_" + list2.back() + "_3_updatedNoveltyMap.ppm");
      if ( beforeComputingNewFeatures )
        destination  =  s_resultdir + NICE::intToString(this->newImageCounter) + "_" + list2.back() + "_0_initialNoveltyMap.ppm";

      noveltyImageRGB.writePPM( destination );
    }
    
    
    // now look where the closest features for the current cluster indices are
    int tmpProtCnt ( 0 );
    for (NICE::VVector::const_iterator protIt = prototypes.begin(); protIt != prototypes.end(); protIt++, tmpProtCnt++)
    {
      double distToNewCluster ( std::numeric_limits<double>::max() );
      int indexOfMostSimFeat( 0 );
      double tmpDist;
      int tmpCnt ( 0 );
      
      for ( NICE::VVector::iterator i = features.begin();
            i != features.end();
            i++, tmpCnt++)
      {
        tmpDist = this->distFunction->calculate( *i, *protIt );
        if ( tmpDist < distToNewCluster )
        {
          distToNewCluster = tmpDist;
          indexOfMostSimFeat = tmpCnt;
        }
      }
      
      int posX ( ( positions[indexOfMostSimFeat] ) [0]  );
      int posY ( ( positions[indexOfMostSimFeat] ) [1]  );
      NICE::Circle circ ( Coord( posX, posY), 2*(tmpProtCnt+1) /* radius*/, Color(200,0,255 ) );
      img.draw(circ);  
    }
    
   if ( b_showResults )
    showImage(img, "Current image and most similar features for current cluster"); 
   else
   {
      std::vector< std::string > list2;
      StringTools::split ( _filename, '/', list2 );      

      std::string destination ( s_resultdir + NICE::intToString(this->newImageCounter-1) + "_" + list2.back() + "_3_updatedCurrentCluster.ppm");
      if ( beforeComputingNewFeatures )
        destination  =  s_resultdir + NICE::intToString(this->newImageCounter) + "_" + list2.back() + "_0_initialCurrentCluster.ppm";
      
      img.writePPM( destination );
   }
   
   return noveltyImageGaussFiltered;
}

NICE::ImageT< int > FeatureLearningPrototypes::evaluateCurrentCodebookByAssignments(const std::string& _filename, const bool& beforeComputingNewFeatures, const bool & _binaryShowLatestPrototype)
{
  std::cerr << "evaluateCurrentCodebookByAssignments" << std::endl;
  NICE::ColorImage img( _filename );
  
  int xsize ( img.width() );
  int ysize ( img.height() );
  
  //variables to store feature information
  NICE::VVector features;
  NICE::VVector cfeatures;
  NICE::VVector positions;

  //compute features
  Globals::setCurrentImgFN ( _filename );
  featureExtractor->extractFeatures ( img, features, positions );
  
  //normalization
  for ( NICE::VVector::iterator i = features.begin();
        i != features.end();
        i++)
  {              
    //normalization :)
    i->normalizeL1();
  }
  
  std::cerr << "normalization done - now look for nearest clusters for every extracted feature" << std::endl;

  NICE::ImageT< int > clusterImage ( xsize, ysize );
  clusterImage.set ( 0 );
   
  NICE::VVector::const_iterator posIt = positions.begin();
  for ( NICE::VVector::const_iterator i = features.begin();
        i != features.end();
        i++, posIt++)
  {              
    
    //loop over codebook representatives
    double minDist ( std::numeric_limits<double>::max() );
    int indexOfNearestCluster ( 0 );
    int clusterCounter ( 0 );
    for (NICE::VVector::const_iterator it =  this->prototypes.begin(); it != this->prototypes.end(); it++, clusterCounter++)
    {
      //compute distance
      double tmpDist ( this->distFunction->calculate(*i,*it) );
      if (tmpDist < minDist)
      {
        minDist = tmpDist;
        indexOfNearestCluster = clusterCounter;
      }
    }

    //take minimum distance and store in in a float image    
    //TODO hard coded!!!
    int noProtoTypes ( this->prototypes.size() -1 );
    
    for ( int tmpY = (*posIt)[1]  - 1; tmpY < (*posIt)[1]  + 1; tmpY++)
    {
      for ( int tmpX = (*posIt)[0]  - 1; tmpX < (*posIt)[0]  + 1; tmpX++)
      {
        if ( _binaryShowLatestPrototype )
        {
          //just a binary image - 1 if newest prototype is nearest - 0 if not
          if  ( indexOfNearestCluster == noProtoTypes)
            clusterImage ( tmpX, tmpY ) = 1;
          else
            clusterImage ( tmpX, tmpY ) = 0;
        }
        else
          //as many different values as current prototypes available
          clusterImage ( tmpX, tmpY ) = indexOfNearestCluster;  
      }
    }
//     clusterImage ( (*posIt)[0], (*posIt)[1] ) = indexOfNearestCluster;
  } 
  
  //show how many clusters we have
  if ( !_binaryShowLatestPrototype )  
  {
    int tmpCnt ( 0 );
    for (NICE::VVector::const_iterator protoIt = prototypes.begin(); protoIt != prototypes.end(); protoIt++, tmpCnt++)
    {
    for ( int tmpY = 1 + 2 - 2; tmpY < (2  + 2); tmpY++)
      {
        for ( int tmpX = 1 + 5*tmpCnt  - 2; tmpX < (1 + 5*tmpCnt  + 2); tmpX++)
        {
          //Take care, this might go "over" the image
          clusterImage ( tmpX, tmpY ) = (Ipp8u) tmpCnt;
        }
      }   
    }
  }
  std::cerr << " evaluateCurrentCodebookByAssignments done" << std::endl;
  return clusterImage;
}