/** 
* @file SemSegSTF.cpp
* @brief Localization system
* @author Erik Rodner
* @date 02/11/2008

*/
#include <iostream>

#include "SemSegSTF.h"
#include "vislearning/baselib/Globals.h"
#include "vislearning/baselib/Preprocess.h"
#include "vislearning/baselib/ProgressBar.h"
#include "core/basics/StringTools.h"
#include "vislearning/baselib/Globals.h"

#include "vislearning/cbaselib/CachedExample.h"
#include "vislearning/cbaselib/PascalResults.h"

#include "vislearning/features/fpfeatures/PixelPairFeature.h"
#include "vislearning/features/fpfeatures/SemanticFeature.h"

#include "vislearning/features/fpfeatures/FIGradients.h"

#include "FIShotton.h"

#include "SemSegTools.h"

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

SemSegSTF::SemSegSTF( const Config *conf, 
		    const MultiDataset *md )
    : SemanticSegmentation ( conf, &(md->getClassNames("train")) )
{
    use_semantic_features = conf->gB("bost", "use_semantic_features", true );
    use_pixelpair_features = conf->gB("bost", "use_pixelpair_features", true );
    subsamplex = conf->gI("bost", "subsamplex", 5);
    subsampley = conf->gI("bost", "subsampley", 5);
    numClasses = md->getClassNames("train").numClasses();

    read_pixel_cache = conf->gB("FPCPixel", "read_cache", false );
    cachepixel = conf->gS("FPCPixel", "cache", "fpc.data" );

    read_seg_cache = conf->gB("FPCSeg", "read_cache", true );
    cacheseg = conf->gS("FPCSeg", "cache", "segforest.data" );

    Examples examples;
    vector<CachedExample *> imgexamples;

    fpcPixel = new FPCRandomForests ( conf, "FPCPixel" );
    fpcPixel->setMaxClassNo ( classNames->getMaxClassno() );
    if ( !read_pixel_cache || !read_seg_cache ) 
    {
	// Generate Positioned Examples
	SemSegTools::collectTrainingExamples ( conf, "FPCPixel", *((*md)["train"]), *classNames, 
	    examples, imgexamples );
    }

    if ( ! read_pixel_cache ) 
    {	
	///////////////////////////////////
	// Train Single Pixel Classifier
	//////////////////////////////////

	FeaturePool fp;
	for ( vector<CachedExample *>::const_iterator k = imgexamples.begin(); 
			k != imgexamples.end();
			k++ )
		fillCachePixel (*k);
    
	PixelPairFeature hf (conf);
	hf.explode ( fp );

	fpcPixel->train ( fp, examples );
	fpcPixel->save ( cachepixel );

	fp.destroy();
    } else {
	fprintf (stderr, "SemSegSTF:: Reading pixel classifier data from %s\n", cachepixel.c_str() );
	fpcPixel->read ( cachepixel );
    }
    
    fpcSeg = new FPCRandomForests ( conf, "FPCSeg" );
    fpcSeg->setMaxClassNo ( classNames->getMaxClassno() );
    maxdepthSegmentationForest = conf->gI("bost", "maxdepth", 5);
    maxdepthSegmentationForestScores = conf->gI("bost", "maxdepth_scores", 9999);

    if ( ! read_seg_cache ) 
    {
	///////////////////////////////////
	// Train Segmentation Forest
	//////////////////////////////////

	fprintf (stderr, "Calculating Prior Statistics\n");
	ProgressBar pbseg ("Calculating Prior Statistics");
	pbseg.show();
	for ( vector<CachedExample *>::const_iterator k = imgexamples.begin(); 
			    k != imgexamples.end();
			    k++ )
	{
	    pbseg.update ( imgexamples.size() );
	    fillCacheSegmentation ( *k );
	}
	pbseg.hide();
	
	FeaturePool fp;

	if ( use_semantic_features )
	{
	    set<int> classnos;
	    classNames->getSelection ( conf->gS("FPCSeg", "train_selection")
		, classnos );
	    SemanticFeature sf ( conf, &classnos );
	    sf.explode ( fp );
	}

	fprintf (stderr, "Training Segmentation Forest\n");

	fpcSeg->train ( fp, examples );
	fpcSeg->save ( cacheseg );

	// clean up memory !!
	for ( vector<CachedExample *>::iterator i = imgexamples.begin();
	    	    i != imgexamples.end();
		    i++ )
	    delete ( *i );

	fp.destroy();
	
    } else {
	fprintf (stderr, "SemSegSTF:: Reading region classifier data from %s\n", cacheseg.c_str() );
	fpcSeg->read ( cacheseg );
	fprintf (stderr, "SemSegSTF:: successfully read\n" );
    }
}

SemSegSTF::~SemSegSTF()
{
}

void SemSegSTF::fillCacheSegmentation ( CachedExample *ce )
{
    FIShotton::buildSemanticMap ( ce,
		      fpcPixel,
		      subsamplex,
		      subsampley, 
		      numClasses );
}

void SemSegSTF::fillCachePixel ( CachedExample *ce )
{
}

void SemSegSTF::semanticseg ( CachedExample *ce, 
			   NICE::Image & segresult,
			   NICE::MultiChannelImageT<double> & probabilities )
{
    int xsize;
    int ysize;
    ce->getImageSize ( xsize, ysize );
    
    int numClasses = classNames->numClasses();
    
    fillCachePixel ( ce );
    fillCacheSegmentation ( ce );

    fprintf (stderr, "BoST classification !\n");

    Example pce ( ce, 0, 0 );

    int xsize_s = xsize / subsamplex;
    int ysize_s = ysize / subsampley;
    ClassificationResult *results = new ClassificationResult [xsize_s*ysize_s];

    /** classify each pixel of the image */
    FullVector prior ( classNames->getMaxClassno() );
    
    probabilities.reInit ( xsize_s, ysize_s, numClasses, true );
    probabilities.setAll ( 0 );

    long offset_s = 0;
    for ( int ys = 0 ; ys < ysize_s ; ys ++ ) 
	for ( int xs = 0 ; xs < xsize_s ; xs++,offset_s++ )  
	{
	    int x = xs * subsamplex;
	    int y = ys * subsampley;
	    pce.x = x ; pce.y = y ; 
	    results[offset_s] = fpcSeg->classify ( pce );

	    for ( int i = 0 ; i < results[offset_s].scores.size(); i++ )
		probabilities.data[i][offset_s] = results[offset_s].scores[i];

	    /*
	    if ( imagePriorMethod != IMAGE_PRIOR_NONE ) 
		prior.add ( results[offset_s].scores );
	    */
	}

   
    fprintf (stderr, "BoST classification ready\n");

    /** save results */
    segresult.resize(xsize_s, ysize_s); 
    segresult.set( classNames->classno("various") );

    long int offset = 0;
    for ( int y = 0 ; y < ysize_s ; y++ ) 
	for ( int x = 0 ; x < xsize_s ; x++,offset++ )  
	{
	    double maxvalue = - numeric_limits<double>::max();
	    int maxindex = 0;
	    for ( int i = 0 ; i < (int)probabilities.numChannels; i++ )
		if ( probabilities.data[i][offset] > maxvalue )
		{
		    maxindex = i;
		    maxvalue = probabilities.data[i][offset];
		}

	    segresult.setPixel(x,y,maxindex);
	}

}