/**
* @file evaluateCompleteBoWPipeline.cpp
* @brief A complete BoW pipeline: feature extraction, codebook creation, vector quantization, classifier training, evaluation on separate test set
* @author Alexander Freytag
* @date 10-05-2013
*/

//STL
#include <iostream>
#include <limits>

//core -- basic stuff
#include <core/basics/Config.h>
#include <core/basics/ResourceStatistics.h>
#include <core/basics/Timer.h>
#include <core/image/Convert.h>
#include <core/vector/VectorT.h>

//vislearning -- basic stuff
#include <vislearning/baselib/Globals.h>
#include <vislearning/baselib/ICETools.h>
#include <vislearning/cbaselib/MultiDataset.h>
#include <vislearning/cbaselib/Example.h>
#include <vislearning/cbaselib/ClassificationResult.h>
#include <vislearning/cbaselib/ClassificationResults.h>
//
// vislearning -- classifier
#include <vislearning/classifier/classifierbase/VecClassifier.h>
#include <vislearning/classifier/genericClassifierSelection.h>
//
// vislearning -- BoW codebooks
#include "vislearning/features/simplefeatures/CodebookPrototypes.h"
#include "vislearning/features/simplefeatures/BoWFeatureConverter.h"
// 
// vislearning -- local features
#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/features/localfeatures/GenericLFSelection.h>
//
// vislearning -- clustering methods
#include <vislearning/math/cluster/ClusterAlgorithm.h>
#include "vislearning/math/cluster/RandomClustering.h"
#include <vislearning/math/cluster/KMeans.h>
#include <vislearning/math/cluster/KMedian.h>
#include <vislearning/math/cluster/GMM.h>
//

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


LocalFeatureRepresentation * setFeatureExtractor( const Config * _conf )
{  
  LocalFeatureRepresentation * featureExtractor;
   
    //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 );   
  
  // Welche Opponentsift Implementierung soll genutzt werden ?
  LocalFeatureRepresentation *cSIFT = NULL;
  LocalFeatureRepresentation *writeFeats = NULL;
  LocalFeatureRepresentation *readFeats = NULL;
  featureExtractor = NULL;
  if ( opSiftImpl == "NICE" )
  {
     cSIFT = new OBJREC::LFonHSG ( _conf, "HSG" );
  }
  else if ( opSiftImpl == "VANDESANDE" )
  {
    cSIFT = new OBJREC::LFColorSande ( _conf, "LFColorSande" );
  }
  else
  {
    fthrow ( Exception, "feattype: %s not yet supported" << opSiftImpl );
  }

  featureExtractor = cSIFT;
  
  if ( writefeat )
  {
    // write the features to a file, if there isn't any to read
    writeFeats = new LFWriteCache ( _conf, cSIFT );
    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 );
    }
    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 ;   
  
  return featureExtractor;
}

OBJREC::ClusterAlgorithm * setClusterAlgo( const Config * _conf )
{ 
 std::string section ( "clusteringStuff" );
  // define the initial number of clusters our codebook shall contain
  int initialNumberOfClusters = _conf->gI(section, "initialNumberOfClusters", 10);
  
  // define the clustering algorithm to be used
  std::string clusterAlgoString = _conf->gS(section, "clusterAlgo", "kmeans");  
  
  OBJREC::ClusterAlgorithm * clusterAlgo;
  
  if (clusterAlgoString.compare("kmeans") == 0)
  {
    clusterAlgo = new OBJREC::KMeans(initialNumberOfClusters);
  }
  else if (clusterAlgoString.compare("kmedian") == 0)
  {
    clusterAlgo = new OBJREC::KMedian(initialNumberOfClusters);
  }  
  else if (clusterAlgoString.compare("GMM") == 0) 
  {
    clusterAlgo = new OBJREC::GMM( _conf, initialNumberOfClusters );
  }
  else if ( clusterAlgoString.compare("RandomClustering") == 0 )   
  {
    clusterAlgo = new OBJREC::RandomClustering( _conf, section );
  }
  else
  {
    std::cerr << "Unknown cluster algorithm selected, use random clustering instead" << std::endl;
    clusterAlgo = new OBJREC::RandomClustering( _conf, section );
  }    
  
  return clusterAlgo;
}


/**
 a complete BoW pipeline
 
 possibly, we can make use of objrec/progs/testClassifier.cpp
*/
int main( int argc, char **argv )
{
  std::set_terminate( __gnu_cxx::__verbose_terminate_handler );

  Config * conf = new Config ( argc, argv );
  
  const bool writeClassificationResults = conf->gB( "main", "writeClassificationResults", true );
  const std::string resultsfile = conf->gS( "main", "resultsfile", "/tmp/results.txt" );
  
  ResourceStatistics rs;
  
  // ========================================================================
  //                            TRAINING STEP
  // ========================================================================

  MultiDataset md( conf );
  const LabeledSet *trainFiles = md["train"];  
   
  //**********************************************
  //
  //     FEATURE EXTRACTION FOR TRAINING IMAGES
  //
  //**********************************************  
  
  std::cerr << "FEATURE EXTRACTION FOR TRAINING IMAGES" << std::endl;
  
  OBJREC::LocalFeatureRepresentation * featureExtractor = OBJREC::GenericLFSelection::selectLocalFeatureRep ( conf, "features", OBJREC::GenericLFSelection::TRAINING );
//   LocalFeatureRepresentation * featureExtractor = setFeatureExtractor( conf );
  
  //collect features in a single data structure
  NICE::VVector featuresFromAllTrainingImages;  
  featuresFromAllTrainingImages.clear();
  
  //okay, this is redundant - but I see no way to do it more easy right now...
  std::vector<NICE::VVector> featuresOfImages ( trainFiles->size() );
  //this again is somehow redundant, but we need the labels lateron for easy access - change this to a better solution :)
  NICE::VectorT<int> labelsTrain ( trainFiles->size(), 0 );
  
  //TODO replace the nasty makro by a suitable for-loop to make it omp-ready (parallelization)
  int imgCnt ( 0 );
  LOOP_ALL_S( *trainFiles )
  {
      EACH_INFO( classno, info );
      std::string filename = info.img();  
      
      NICE::ColorImage img( filename );
  
      //compute features
      
      //variables to store feature information
      NICE::VVector features;
      NICE::VVector positions;
  
      Globals::setCurrentImgFN ( filename );
      featureExtractor->extractFeatures ( img, features, positions );  
        
      //normalization :)
      for ( NICE::VVector::iterator i = features.begin();
            i != features.end();
            i++)
      {              
        i->normalizeL1();
      }  
      
      //collect them all in a larger data structure
      featuresFromAllTrainingImages.append( features );
      //and store it as well in the data struct that additionally keeps the information which features belong to which image
      //TODO this can be made more clever!
//       featuresOfImages.push_back( features );
      featuresOfImages[imgCnt] = features;
      labelsTrain[imgCnt] = classno;
      imgCnt++;
  }
  
  
  
  //**********************************************
  //
  //             CODEBOOK CREATION
  //
  //**********************************************    
  
  std::cerr << "CODEBOOK CREATION" << std::endl;
  OBJREC::ClusterAlgorithm * clusterAlgo = setClusterAlgo( conf );
   
  NICE::VVector prototypes;
  
  std::vector<double> weights;
  std::vector<int> assignments;
  
  clusterAlgo->cluster( featuresFromAllTrainingImages, prototypes, weights, assignments );
  
  
  OBJREC::CodebookPrototypes * codebook = new OBJREC::CodebookPrototypes ( prototypes );
  
  
  //**********************************************
  //
  //             VECTOR QUANTIZATION OF
  //           FEATURES OF TRAINING IMAGES
  //
  //**********************************************  
  
  OBJREC::BoWFeatureConverter * bowConverter = new OBJREC::BoWFeatureConverter ( conf, codebook );  
  
  OBJREC::LabeledSetVector trainSet;
  
  NICE::VVector histograms ( featuresOfImages.size() /* number of vectors*/, 0 /* dimension of vectors*/ ); //the internal vectors will be resized within calcHistogram
  NICE::VVector::iterator histogramIt = histograms.begin();
  NICE::VectorT<int>::const_iterator labelsIt = labelsTrain.begin();
  
  for (std::vector<NICE::VVector>::const_iterator imgIt = featuresOfImages.begin(); imgIt != featuresOfImages.end(); imgIt++, histogramIt++, labelsIt++)
  {
    bowConverter->calcHistogram ( *imgIt, *histogramIt );
    bowConverter->normalizeHistogram ( *histogramIt );
    
    //NOTE perhaps we should use add_reference here
    trainSet.add( *labelsIt, *histogramIt );
  }
  
  
  //**********************************************
  //
  //             CLASSIFIER TRAINING
  //
  //**********************************************  
  
  std::string classifierType = conf->gS( "main", "classifierType", "GPHIK" );
  OBJREC::VecClassifier * classifier = OBJREC::GenericClassifierSelection::selectVecClassifier( conf, classifierType );
  
  //TODO integrate GP-HIK-NICE into vislearning and add it into genericClassifierSelection
  
  //this method adds the training data to the temporary knowledge of our classifier
  classifier->teach( trainSet );
  //now the actual training step starts (e.g., parameter estimation, ... )
  classifier->finishTeaching();
  
  
  // ========================================================================
  //                            TEST STEP
  // ========================================================================
  
  const LabeledSet *testFiles = md["test"];
  
  NICE::Matrix confusionMat;
  NICE::Timer t;
  
  ClassificationResults results;
  
  LOOP_ALL_S( *testFiles )
  {
      EACH_INFO( classno, info );
      std::string filename = info.img();
  
      //**********************************************
      //
      //     FEATURE EXTRACTION FOR TEST IMAGES
      //
      //**********************************************  
      
      NICE::ColorImage img( filename );
  
      //compute features
      
      //variables to store feature information
      NICE::VVector features;
      NICE::VVector positions;
  
      Globals::setCurrentImgFN ( filename );
      featureExtractor->extractFeatures ( img, features, positions );  
        
      //normalization :)
      for ( NICE::VVector::iterator i = features.begin();
            i != features.end();
            i++)
      {              
        i->normalizeL1();
      }        
      
      //**********************************************
      //
      //             VECTOR QUANTIZATION OF
      //           FEATURES OF TEST IMAGES
      //
      //********************************************** 
      
      NICE::Vector histogramOfCurrentImg;
      bowConverter->calcHistogram ( features, histogramOfCurrentImg );
      bowConverter->normalizeHistogram ( histogramOfCurrentImg );
      
      //**********************************************
      //
      //             CLASSIFIER EVALUATION
      //
      //**********************************************   
            
      uint classno_groundtruth = classno;
      
      t.start();
      ClassificationResult r = classifier->classify ( histogramOfCurrentImg );
      t.stop();
      uint classno_estimated = r.classno;

      //if we like to store the classification results for external post processing, uncomment this
      if ( writeClassificationResults )
      {
        results.push_back( r );
      }      
      
      confusionMat( classno_estimated, classno_groundtruth ) += 1;
  }
  
  confusionMat.normalizeColumnsL1();
  std::cerr << confusionMat << std::endl;

  std::cerr << "average recognition rate: " << confusionMat.trace()/confusionMat.rows() << std::endl;  
  
  if ( writeClassificationResults )
  {
    double avgRecogResults  ( results.getAverageRecognitionRate () );    
    std::cerr << "average recognition rate according to classificationResults: " << avgRecogResults << std::endl;
    results.writeWEKA ( resultsfile, 0 );
  }     
  
 
  
   return 0;
}