#include "FeatureLearningPrototypes.h"

//STL
#include <iostream>

//core
#include <core/image/FilterT.h>
#include <core/image/CircleT.h>
#include <core/image/Convert.h>
#include <core/imagedisplay/ImageDisplay.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/KMedian.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("kmedian") == 0)
  {
    this->clusterAlgo = new OBJREC::KMedian(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 from " << cacheInitialCodebook << 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  );
    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 distToCurrentCluster ( std::numeric_limits<double>::max() );
      int indexOfMostSimFeat( 0 );
      double tmpDist;
      int featureCnt ( 0 );
      
      for ( NICE::VVector::iterator i = features.begin();
            i != features.end();
            i++, featureCnt++)
      {
        tmpDist = this->distFunction->calculate( *i, *protIt );
        if ( tmpDist < distToCurrentCluster )
        {
          distToCurrentCluster = tmpDist;
          indexOfMostSimFeat = featureCnt;
        }
      }
      
      int posX ( ( positions[indexOfMostSimFeat] ) [0]  );
      int posY ( ( positions[indexOfMostSimFeat] ) [1]  );
      
      //position (for OpponentSIFT of van de Sande): x y scale orientation cornerness
      
      /*What is the interpretation of scale?

      The scale parameter was implemented to correspond with the Gaussian filter sigma at which points were detected. Therefore, the
      scale is not directly interpretable as the size of the region described in terms of number of pixels. However, it is linearly
      related the radius of the circular area described. To capture the area of the Gaussian originally used, we have a 3x 
      magnification factor. But, remember that SIFT has 4x4 cells, and this is a measure for a single cell. So, because we are 
      working with a radius, multiply by 2. Due to the square shape of the region, we need to extend the outer radius even further 
      by sqrt(2), otherwise the corners of the outer cells are cut off by our circle. So, the largest outer radius is 
      Round(scale * 3 * 2 * sqrt(2)). The area from which the SIFT descriptor is computed is a square which fits within a circle 
      of this radius. Also, due to the Gaussian weighting applied within SIFT, the area that really matters is much, much smaller: 
      the outer parts get a low weight. 

      For the default scale of 1.2, we get a outer circle radius of 10. The potential sampling area then becomes -10..10, e.g. a 
      21x21patch. However, the square area which fits inside this circle is smaller: about 15x15. The corners of this 15x15square 
      touch the outer circle. */
      
      /*Why is the direction (angle) field always 0?


      Estimating the dominant orientation of a descriptor is useful in matching scenarios. However, in an object/scene categorization
      setting, the additional invariance reduces accuracy. Being able to discriminate between dominant directions of up and right 
      is very useful here, and rotated down images are quite rare in an object categorization setting. Therefore, orientation 
      estimation is disabled in the color descriptor software. The subsequent rotation of the descriptor to achieve 
      rotation-invariance is still possible by supplying your own regions and angles for an image (through --loadRegions). However, 
      by default, no such rotation is performed, since the default angle is 0. */      
      
      
      //adapt the pseudo color transformation as done in Convert.cpp
        size_t seg = ( size_t ) ( tmpProtCnt/(float)prototypes.size() * 6.0 );
        double y   = ( 6 * tmpProtCnt/(float)prototypes.size() - seg );

        Color circleColor;
        switch ( seg ) {
          case 0:
            circleColor = Color( 0,0,(int)round(y*255) );
            break;
          case 1:
            circleColor = Color( 0,(int)round(y*255),255 );
            break;
          case 2:
            circleColor = Color( 0,255,(int)round((1-y)*255) );
            break;
          case 3:
            circleColor = Color( (int)round(y*255),255,0 );
            break;
          case 4:
            circleColor = Color( 255,(int)round((1-y)*255),0 );
            break;
          case 5:
            circleColor = Color( 255,(int)round(y*255),(int)round(y*255) );
            break;
          default:
            circleColor = Color( 255,255,255 );
            break;
        }      
      
      NICE::Circle circ ( Coord( posX, posY), (int) round(2*3*sqrt(2)*( positions[indexOfMostSimFeat] )[2]) /* radius*/, circleColor );
      img.draw(circ);  
    }
        
   if ( b_showResults )
      showImage(img, "Current image and most similar features for current prototypes"); 
   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();
  }
  
  //this is the image we will return finally
  NICE::ImageT< int > clusterImage ( xsize, ysize );
  clusterImage.set ( 0 );
  
  // after iterating over all features from the new image, this vector contains
  // distances to the most similar feature for every prototype
  NICE::Vector minDistances ( this->prototypes.size() );
  
  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 indexOfNearestPrototype ( 0 );
    int prototypeCounter ( 0 );
    for (NICE::VVector::const_iterator it =  this->prototypes.begin(); it != this->prototypes.end(); it++, prototypeCounter++)
    {
      //compute distance
      double tmpDist ( this->distFunction->calculate(*i,*it) );
      //check what the closest prototype is
      if (tmpDist < minDist)
      {
        minDist = tmpDist;
        indexOfNearestPrototype = prototypeCounter;
      }
      
      //check whether we found a feature for the current prototype which is more similar then the previous best one
      if ( minDistances[ prototypeCounter ] > tmpDist )
        minDistances[ prototypeCounter ] = tmpDist;      
    }
    

    
    
    //take minimum distance and store in in a float image    
    // for nice visualization, we plot the cluster index into a square of size 3 x 3
    //TODO currently 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  ( indexOfNearestPrototype == noProtoTypes)
            clusterImage ( tmpX, tmpY ) = 1;
          else
            clusterImage ( tmpX, tmpY ) = 0;
        }
        else
          //as many different values as current prototypes available
          clusterImage ( tmpX, tmpY ) = indexOfNearestPrototype;  
      }
    }
  } 
  
  std::cerr << "Codebook evaluation by assignments... min distances in image for every prototype: " << std::endl << "    " << minDistances << std::endl;
  
  //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 + 4*tmpCnt ; tmpX < (1 + 4*tmpCnt  + 3); tmpX++)
          {
            //Take care, this might go "over" the image
            clusterImage ( tmpX, tmpY ) = (Ipp8u) tmpCnt;
          }
        }   
    }
  }
  
  return clusterImage;
}

void FeatureLearningPrototypes::evaluateCurrentCodebookByConfusionMatrix( NICE::Matrix & _confusionMat )
{
  _confusionMat.resize ( this->prototypes.size(), this->prototypes.size() );
  _confusionMat.set( 0.0 );
  
  double tmpDist ( 0.0 );
  NICE::VVector::const_iterator protoItJ = prototypes.begin();
  for ( int j = 0; j < prototypes.size(); j++, protoItJ++)
  {
    NICE::VVector::const_iterator protoItI = protoItJ; 
    for ( int i = j; i < prototypes.size(); i++, protoItI++)
    {
      tmpDist = this->distFunction->calculate( *protoItJ, *protoItI );
      
      //assuming symmetric distances
      _confusionMat ( i, j ) = tmpDist;
//       if ( i != j )
//         _confusionMat ( j, i ) = tmpDist;  
    }
  }
}

NICE::VVector * FeatureLearningPrototypes::getCurrentCodebook()
{
  return &(this->prototypes);
}