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

*/

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

using namespace OBJREC;

using namespace std;

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::ImageT<int> & result, NICE::MultiChannelImageT<double> & probabilities )
{
  assert ( prior.size() == ( int ) probabilities.channels() );
  int xsize = probabilities.width();
  int ysize = probabilities.height();

  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 );

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

          if ( probabilities[i](xs,ys) > maxvalue )
          {
            maxindex = i;
            maxvalue = probabilities[i](xs,ys);
          }
        }

        if ( sum > 1e-11 )
          for ( int i = 0 ; i < ( int ) probabilities.channels() ; i++ )
          {
            probabilities[i](xs,ys) /= sum;
          }

        result.setPixel ( xs, ys, maxindex );
      }
  } else if ( imagePriorMethod == IMAGE_PRIOR_PSEUDOPROB ) {
    for ( int ys = 0 ; ys < ysize ; ys ++ )
      for ( int xs = 0 ; xs < xsize ; xs++)
      {
        int maxindex = 0;
        double maxvalue = - numeric_limits<double>::max();
        double sum = 0.0;

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