#include "FeatureLearningClusterBased.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/baselib/ICETools.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 FeatureLearningClusterBased::setClusterAlgo( const std::string & _clusterAlgoString, const bool & _setForInitialTraining)
{
  //be careful with previously allocated memory
  if (this->clusterAlgo != NULL)
    delete clusterAlgo;
  
  if (_clusterAlgoString.compare("kmeans") == 0)
  {
    if ( _setForInitialTraining )
      this->clusterAlgo = new OBJREC::KMeans(this->initialNumberOfClusters);
    else
      this->clusterAlgo = new OBJREC::KMeans(this->numberOfClustersForNewImage);
  }
  else if (_clusterAlgoString.compare("GMM") == 0) 
  {
    if ( _setForInitialTraining )
      this->clusterAlgo = new OBJREC::GMM(this->conf, this->initialNumberOfClusters);
    else
      this->clusterAlgo = new OBJREC::GMM(this->conf, this->numberOfClustersForNewImage);      
  }
  else
  {
    std::cerr << "Unknown cluster algorithm selected, use k-means instead" << std::endl;
    if ( _setForInitialTraining )      
      this->clusterAlgo = new OBJREC::KMeans(this->initialNumberOfClusters);
    else
      this->clusterAlgo = new OBJREC::KMeans(this->numberOfClustersForNewImage);
  }    
}

void FeatureLearningClusterBased::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 ( 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 FeatureLearningClusterBased::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 FeatureLearningClusterBased::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 FeatureLearningClusterBased::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
  //
  //********************************************** 


FeatureLearningClusterBased::FeatureLearningClusterBased ( 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 ); 
  
  showTrainingImages = conf->gB( section, "showTrainingImages", false );  
  showResults = conf->gB( section, "showResults", false );
  
  resultdir = conf->gS( section, "resultdir", "/tmp/");
  
  
  //! define the initial number of clusters our codebook shall contain
  initialNumberOfClusters = conf->gI(section, "initialNumberOfClusters", 10);
   //! define the number of clusters we want to compute for an unseen image
  numberOfClustersForNewImage = conf->gI(section, "numberOfClustersForNewImage", 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, true /*set cluster algo for training*/ );
  
  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 ;   
  
  this->setClusterAlgo( clusterAlgoString, false /*set cluster algo for feature learning*/ );
  
  //so far, we have not seen any new image
  this->newImageCounter = 0;
  
  //TODO stupid
  this->maxValForVisualization = 0.005;
}

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

void FeatureLearningClusterBased::learnNewFeatures ( const std::string & _filename )
{  
  NICE::ColorImage img( _filename );
  
  int xsize ( img.width() );
  int ysize ( img.height() );
   
  //variables to store feature information
  NICE::VVector newFeatures;
  NICE::VVector cfeatures;
  NICE::VVector positions;

  //compute features
  std::cerr << " EXTRACT FEATURES FROM UNSEEN IMAGE" << std::endl;
  Globals::setCurrentImgFN ( _filename );
  featureExtractor->extractFeatures ( img, newFeatures, positions );  
  
  //store feature information in larger data structure
  for ( NICE::VVector::iterator i = newFeatures.begin();
        i != newFeatures.end();
        i++)
  {              
    //normalization :)
    i->normalizeL1();
  }
  
  //cluster features
  std::cerr << " CLUSTER FEATURES FROM UNSEEN IMAGE" << std::endl;
  NICE::VVector prototypesForNewImage;
  std::vector< double > weights;
  std::vector< int > assignment;
  clusterAlgo->cluster ( newFeatures, prototypesForNewImage, weights, assignment);
  
  if ( b_evaluationWhileFeatureLearning )
  {
    //visualize new clusters
    int tmpProtCnt ( 0 );    
    for (NICE::VVector::const_iterator protIt = prototypesForNewImage.begin(); protIt != prototypesForNewImage.end(); protIt++, tmpProtCnt++)
    {
      double distToNewCluster ( std::numeric_limits<double>::max() );
      int indexOfMostSimFeat( 0 );
      double tmpDist;
      int tmpCnt ( 0 );
      
      for ( NICE::VVector::iterator i = newFeatures.begin();
            i != newFeatures.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), 10 /* radius*/, Color(200,0,255) );
      img.draw(circ);       
    }
    
    
    //draw features most similar to old clusters
    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 = newFeatures.begin();
            i != newFeatures.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), 5 /* radius*/, Color(200,255,0 ) );
      img.draw(circ);      
    }
    
    if ( showResults )
      showImage(img, "Current (new) image and most similar feature for new cluster");     
    else
    {
      std::vector< std::string > list2;
      StringTools::split ( _filename, '/', list2 );      

      std::string destination ( resultdir + NICE::intToString(this->newImageCounter) + "_" + list2.back() + "_1_oldAndNewClusters.ppm");
      img.writePPM( destination );
    }   
  }
  
  //compute score for every cluster: #assigned features * distance to current cluster centers
  
  NICE::Vector distancesToCurrentClusters ( numberOfClustersForNewImage, 0.0 );
  NICE::Vector clusterSizes ( numberOfClustersForNewImage, 0.0 ); //i.e., the number of assignments, or a derived number
   
  //compute "relevance" of every new cluster
 
  std::cerr << " COMPUTE SIZES OF NEW CLUSTERS" << std::endl;
  for (std::vector<int>::const_iterator assignIt = assignment.begin(); assignIt != assignment.end(); assignIt++)
  {
    clusterSizes[*assignIt]++;
  }
  clusterSizes.normalizeL1();
  
  std::cerr << "cluster Sizes: " << clusterSizes << std::endl;
  
  
  //compute distances of new cluster centers to old cluster centers
  std::cerr << " COMPUTE DISTANCES BETWEEN NEW AND OLD CLUSTERS" << std::endl;
  NICE::Vector::iterator distanceIt = distancesToCurrentClusters.begin();
  for ( NICE::VVector::const_iterator newProtIt = prototypesForNewImage.begin(); newProtIt != prototypesForNewImage.end(); newProtIt++, distanceIt++)
  {
    double minDist ( std::numeric_limits<double>::max() );
    double tmpDist;
    for ( NICE::VVector::const_iterator protIt = prototypes.begin(); protIt != prototypes.end(); protIt ++)
    {
        //compute distance
        tmpDist = this->distFunction->calculate( *protIt, *newProtIt );
        if (tmpDist < minDist)
          minDist = tmpDist;      
    }
    
    *distanceIt = minDist;
  }
  
  std::cerr << "distances: " << distancesToCurrentClusters << std::endl;
  
  //compute final scores for the new image
  NICE::Vector clusterScores ( numberOfClustersForNewImage, 0.0 );
  for (uint i = 0; i < numberOfClustersForNewImage; i++)
  {
      clusterScores[i] = clusterSizes[i] * distancesToCurrentClusters[i];
  }
  
  std::cerr << "final cluster scores for new image: " << clusterScores << std::endl;
    
  NICE::Vector chosenClusterCenter ( prototypesForNewImage[ clusterScores.MaxIndex()  ] );
  
  
  //include the chosen information into the currently used prototypes
  prototypes.push_back( chosenClusterCenter );
  
  if ( b_evaluationWhileFeatureLearning ) 
  {
    
    NICE::ColorImage imgTmp( _filename );
    
    double distToNewCluster ( std::numeric_limits<double>::max() );
    int indexOfMostSimFeat( 0 );
    double tmpDist;
    int tmpCnt ( 0 );
    
    for ( NICE::VVector::iterator i = newFeatures.begin();
          i != newFeatures.end();
          i++, tmpCnt++)
    {
      tmpDist = this->distFunction->calculate( *i, chosenClusterCenter );
      if ( tmpDist < distToNewCluster )
      {
        distToNewCluster = tmpDist;
        indexOfMostSimFeat = tmpCnt;
      }
    }
    
    int posX ( ( positions[indexOfMostSimFeat] ) [0]  );
    int posY ( ( positions[indexOfMostSimFeat] ) [1]  );
    NICE::Circle circ ( Coord( posX, posY), 10 /* radius*/, Color(200,0,255) );
    imgTmp.draw(circ); 
    
    if ( showResults )
      showImage(imgTmp, "Current (new) image and most similar feature for new cluster"); 
    else 
    {
      std::vector< std::string > list2;
      StringTools::split ( _filename, '/', list2 );      

      std::string destination ( resultdir + NICE::intToString(this->newImageCounter) + "_" + list2.back() + "_2_bestNewCluster.ppm");
      imgTmp.writePPM( destination );
    }
  }
  
  //this was a new image, so we increase our internal counter
  (this->newImageCounter)++;  
}

NICE::FloatImage FeatureLearningClusterBased::evaluateCurrentCodebook ( const std::string & _filename , const bool & beforeComputingNewFeatures )
{
    NICE::ColorImage img( _filename );
    if ( 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 ( showResults )
      showImage(noveltyImageRGB, "Novelty Image");  
    else 
    {
      std::vector< std::string > list2;
      StringTools::split ( _filename, '/', list2 );      

      std::string destination ( resultdir + NICE::intToString(this->newImageCounter -1 ) + "_" + list2.back() + "_3_updatedNoveltyMap.ppm");
      if ( beforeComputingNewFeatures )
        destination  =  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 ( 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 ( resultdir + NICE::intToString(this->newImageCounter-1) + "_" + list2.back() + "_3_updatedCurrentCluster.ppm");
      if ( beforeComputingNewFeatures )
        destination  =  resultdir + NICE::intToString(this->newImageCounter) + "_" + list2.back() + "_0_initialCurrentCluster.ppm";
      
      img.writePPM( destination );
   }
   
   return noveltyImageGaussFiltered;
}