#include "FeatureLearningClusterBased.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/baselib/Globals.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);
  }    
}


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


FeatureLearningClusterBased::FeatureLearningClusterBased ( const Config *_conf,
                               const MultiDataset *_md, const std::string & _section )
    : FeatureLearningPrototypes ( _conf, _md, _section )
{ 
  
   // define the number of clusters we want to compute for an unseen image
  numberOfClustersForNewImage = conf->gI(section, "numberOfClustersForNewImage", 10);
   
  //**********************************************
  //
  //      SET UP VARIABLES AND METHODS
  //             - FEATURE TYPE
  //             - CLUSTERING ALGO
  //             - DISTANCE FUNCTION
  //             - ...
  //
  //**********************************************  
  
   
  //run the training to initially compute a codebook and stuff like that
//  this->train( _md );
   
  // define the clustering algorithm to be used
  std::string clusterAlgoString = conf->gS(section, "clusterAlgo", "kmeans");  
  this->setClusterAlgo( clusterAlgoString, false /*set cluster algo for feature learning*/ );
}

FeatureLearningClusterBased::~FeatureLearningClusterBased()
{
  // clean-up
}

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 );  
  
  //normalization :)
  for ( NICE::VVector::iterator i = newFeatures.begin();
        i != newFeatures.end();
        i++)
  {              
    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 ( b_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 ( s_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 ( b_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 ( s_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)++;  
}