#include "FeatureLearningPrototypes.h"

#include <iostream>

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

#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::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();  
  }
  
  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;
}


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


FeatureLearningPrototypes::FeatureLearningPrototypes ( const Config *_conf,
                               const MultiDataset *_md, const std::string & _section )
    : FeatureLearningGeneric ( _conf, _section )
{ 
    //feature stuff
  // which OpponentSIFT implementation to use {NICE, VANDESANDE}
  std::string opSiftImpl;  
  opSiftImpl = conf->gS ( "Descriptor", "implementation", "VANDESANDE" );
  // read features?
  bool readfeat;
  readfeat = conf->gB ( "Descriptor", "read", true );
  // write features?
  bool writefeat;  
  writefeat = conf->gB ( "Descriptor", "write", true ); 
 
  
  // 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;
  
  // Welche Opponentsift Implementierung soll genutzt werden ?
  LocalFeatureRepresentation *cSIFT = NULL;
  LocalFeatureRepresentation *writeFeats = NULL;
  LocalFeatureRepresentation *readFeats = NULL;
  this->featureExtractor = NULL;
  if ( opSiftImpl == "NICE" )
  {
    cSIFT = new OBJREC::LFonHSG ( conf, "HSGtrain" );
  }
  else if ( opSiftImpl == "VANDESANDE" )
  {
    cSIFT = new OBJREC::LFColorSande ( conf, "LFColorSandeTrain" );
  }
  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 ( conf, cSIFT );
    this->featureExtractor = writeFeats;
  }

  if ( readfeat )
  {
    // read the features from a file
    if ( writefeat )
    {
      readFeats = new LFReadCache ( conf, writeFeats, -1 );
    }
    else
    {
      readFeats = new LFReadCache ( conf, cSIFT, -1 );
    }
    this->featureExtractor = readFeats; 
  }
  
  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 );
  
  //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 ;   
    
  //so far, we have not seen any new image
  this->newImageCounter = 0;
  
  //TODO stupid
  this->maxValForVisualization = 0.005;
}

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

NICE::FloatImage FeatureLearningPrototypes::evaluateCurrentCodebook ( const std::string & _filename , const bool & beforeComputingNewFeatures )
{
    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 );
    
    FloatImage noveltyImage ( xsize, ysize );
    noveltyImage.set ( 0.0 );        
    
    double maxDist ( 0.0 );
    
    NICE::VVector::const_iterator posIt = positions.begin();
    //store feature information in larger data structure
    for ( NICE::VVector::iterator i = features.begin();
          i != features.end();
          i++, posIt++)
    {              
      //normalization :)
      i->normalizeL1();
      
      //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;
      }
      
      if (minDist > maxDist)
        maxDist = minDist;

      //take minimum distance and store in in a float image
      
      noveltyImage ( (*posIt)[0], (*posIt)[1]  ) = minDist;
    }      
    
    //gauss-filtering for nicer visualization
    FloatImage noveltyImageGaussFiltered ( xsize, ysize );
    float sigma ( 3.0 );
    FilterT<float, float, float> filter;
    filter.filterGaussSigmaApproximate ( noveltyImage, sigma, &noveltyImageGaussFiltered );
    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;
    //for suitable visualization of scores between zero (known) and one (unknown)
//       noveltyImageGaussFiltered( 0 , 0 ) = std::max<double>(maxDist, 1.0);

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