/** 
* @file PSSImageLevelPrior.cpp
* @brief incorporate prior from image categorization method
* @author Erik Rodner
* @date 03/19/2009

*/
#ifdef NOVISUAL
#include <objrec/nice_nonvis.h>
#else
#include <objrec/nice.h>
#endif

#include <iostream>
#include <set>
#include <assert.h>
#include <algorithm>
#include "PSSImageLevelPrior.h"

using namespace OBJREC;

using namespace std;
// refactor-nice.pl: check this substitution
// old: using namespace ice;
using namespace NICE;



PSSImageLevelPrior::PSSImageLevelPrior( int imagePriorMethod, int priorK, double alphaImagePrior )
{
    this->imagePriorMethod = imagePriorMethod;
    this->priorK = priorK;
    this->alphaImagePrior = alphaImagePrior;
}

PSSImageLevelPrior::~PSSImageLevelPrior()
{
}

void PSSImageLevelPrior::setPrior ( FullVector & prior )
{
    this->prior = prior;
}

void PSSImageLevelPrior::postprocess ( NICE::Image & result, GenericImage<double> & probabilities )
{
    assert ( prior.size() == (int)probabilities.numChannels );
    int xsize = probabilities.xsize;
    int ysize = probabilities.ysize;

    if ( imagePriorMethod == IMAGE_PRIOR_BEST_K ) 
    {
	vector<int> indices;
	prior.getSortedIndices ( indices );

	reverse ( indices.begin(), indices.end() );
	set<int> bestComponents;
	vector<int>::const_iterator j = indices.begin();
	if ( indices.size() > (size_t)priorK )
	    advance ( j, priorK );
	else
	    j = indices.end();

	for ( vector<int>::const_iterator jj = indices.begin();
			    jj != j ; jj++ )
	    bestComponents.insert ( *jj );
	
	int offset_s = 0;

	for (  int ys = 0 ; ys < ysize ; ys ++ ) 
	    for ( int xs = 0 ; xs < xsize ; xs++,offset_s++ )  
	    {
		int maxindex = 0;
		double maxvalue = - numeric_limits<double>::max();
		double sum = 0.0;
		for ( int i = 0 ; i < (int)probabilities.numChannels ; i++ )
		{
		    if ( bestComponents.find(i) == bestComponents.end()  )
			probabilities.data[i][offset_s] = 0.0;
		    sum += probabilities.data[i][offset_s];

		    if ( probabilities.data[i][offset_s] > maxvalue )
		    {
			maxindex = i;
			maxvalue = probabilities.data[i][offset_s];
		    }
		}
	    	
		if ( sum > 1e-11 )
		    for ( int i = 0 ; i < (int)probabilities.numChannels ; i++ )
		    {
			probabilities.data[i][offset_s] /= sum;
		    }
		
		result.setPixel(xs,ys,maxindex);
	    }
    } else if ( imagePriorMethod == IMAGE_PRIOR_PSEUDOPROB ) {
	int offset_s = 0;
	for (  int ys = 0 ; ys < ysize ; ys ++ ) 
	    for ( int xs = 0 ; xs < xsize ; xs++,offset_s++ )  
	    {
		int maxindex = 0;
		double maxvalue = - numeric_limits<double>::max();
		double sum = 0.0;

		for ( int i = 0 ; i < (int)probabilities.numChannels ; i++ )
		{
		    probabilities.data[i][offset_s] *= pow ( prior[i], alphaImagePrior );
		    sum += probabilities.data[i][offset_s];
		    if ( probabilities.data[i][offset_s] > maxvalue )
                    {
                        maxindex = i;
                        maxvalue = probabilities.data[i][offset_s];
		    }										                    
		}
		if ( sum > 1e-11 )
		    for ( int i = 0 ; i < (int)probabilities.numChannels ; i++ )
		    {
			probabilities.data[i][offset_s] /= sum;
		    }
		result.setPixel(xs,ys,maxindex);
	    }
    }
}