/*!
\file skinmodel.cpp

\author Juraj Blaho

\brief Skin model.


*/

#include <cfloat>

#include <string>
#include <fstream>

#include "cvext.hpp"
#include "BlahoException.h"

#include "BlahoSkinmodel.h"

/*!
\brief Convert dimension value to the bin number.
*/
#define VALUE_TO_BIN(VALUE, BIN_COUNT) cvFloor((VALUE)*((BIN_COUNT)-1))

/*!
\brief Conver bin number to the dimension value.
*/
#define BIN_TO_VALUE(BIN, BIN_COUNT) (((double)(BIN))/((BIN_COUNT)-1))

/*!
\brief Cycle through all histogram bins, current index is in the i[] array
*/
#define FOR_ALL_HIST_BINS \
int h_index[COLOR_DIMENSIONS]={0};\
for(;h_index[0]<_hist_resolution[0];h_index[0]++)\
	for(h_index[1]=0;h_index[1]<_hist_resolution[1];h_index[1]++)
	
#define HIST_INDEX h_index[0]][h_index[1]

bool SkinModel::TransformColor(int r, int g, int b, double *output){	
	int i=r+g+b;

	if(i==0)return false;
	
	output[0]=((double)r)/i;
	output[1]=((double)g)/i;
	
	return true;
}

/*
double SkinModel::NormalDistribution(double x, double mean, double var){
	return exp(-(x-mean)*(x-mean)/(2*var))/sqrt(var*6.283185);
}
*/
double SkinModel::NormalDistribution2D(double const *x, double const *mean, double const *var){
	return exp(-(x[0]-mean[0])*(x[0]-mean[0])/(2*var[0])-(x[1]-mean[1])*(x[1]-mean[1])/(2*var[1]))/
	       (sqrt(var[0]*var[1])*6.283185);
}

void SkinModel::ProbabilityUpdateRetrained(){
	double ratio_skin=_skin_prior/(_hist_resolution[0]*_hist_resolution[1]);
	double ratio_bg=(1.0-_skin_prior);
	
	int h_index[ COLOR_DIMENSIONS] = {0};

	for( ; h_index[ 0] < _hist_resolution[ 0]; h_index[ 0]++){
		for( h_index[ 1]=0; h_index[ 1] < _hist_resolution[ 1]; h_index[ 1]++){
			double x[ COLOR_DIMENSIONS];
			double value;
			
			for(int k=0;k<COLOR_DIMENSIONS;k++)
				x[k]=BIN_TO_VALUE(h_index[k],_hist_resolution[k]);
			
			value=NormalDistribution2D(x,_skin_mean,_skin_var)*ratio_skin;
			value/=value+.1;//(_bg[HIST_INDEX])*ratio_bg;
			
			_p_diff[HIST_INDEX]=value;
		}
	}
}

void SkinModel::ProbabilityUpdateAdapted(double tau){
	double mean[COLOR_DIMENSIONS];
	double var[COLOR_DIMENSIONS];
	
	/*calculate adapted skin params*/
	for(int i=0;i<COLOR_DIMENSIONS;i++){
		mean[i]=(_skin_mean[i]*tau+_acc_skin_sum[i])/(tau+_acc_skin_count);
		var[i]=_acc_skin_sqrsum[i]-_acc_skin_sum[i]*_acc_skin_sum[i]/(_acc_skin_count);
		var[i]=(var[i]+tau*_skin_var[i])/(tau+_acc_skin_count);
	}
	
	double ratio_skin=_skin_prior/(_hist_resolution[0]*_hist_resolution[1]);
	double ratio_bg=(1.0-_skin_prior)/(tau+_acc_bg_count);
	
	FOR_ALL_HIST_BINS{
		double x[COLOR_DIMENSIONS];
		double value;
		
		for(int k=0;k<COLOR_DIMENSIONS;k++)
			x[k]=BIN_TO_VALUE(h_index[k],_hist_resolution[k]);
		
		value=NormalDistribution2D(x,mean,var)*ratio_skin;
		value/=value+0.002;//(_bg[HIST_INDEX]*tau+_acc_bg[HIST_INDEX])*ratio_bg;
		
		_p_diff[HIST_INDEX]=value;
	}
		
}

void SkinModel::Init(int r_resolution, int g_resolution){
	_hist_resolution[0]=r_resolution;
	_hist_resolution[1]=g_resolution;
	//_hist_bin_count=((_hist_resolution[0]+1)*_hist_resolution[1]+1)/2;
	_hist_bin_count=_hist_resolution[0]*_hist_resolution[1];
	
	double default_value=1.0/_hist_bin_count;
	
	_skin_prior=0.5;
	
	_p_diff=new double*[_hist_resolution[0]];
	_bg=new double*[_hist_resolution[0]];
	_acc_bg=new double*[_hist_resolution[0]];
	
	for(int i=0;i<_hist_resolution[0];i++){
		_p_diff[i]=new double[_hist_resolution[1]];
		_bg[i]=new double[_hist_resolution[1]];
		_acc_bg[i]=new double[_hist_resolution[1]];
		
		for(int j=0;j<_hist_resolution[1];j++){
			_p_diff[i][j]=0.0;
			_bg[i][j]=default_value;
			_acc_bg[i][j]=0.0;
		}
	}
	
	for(int i=0;i<COLOR_DIMENSIONS;i++){
		_skin_mean[i]=0.0;
		_skin_var[i]=0.0;
		_acc_skin_sum[i]=0.0;
		_acc_skin_sqrsum[i]=0.0;
	}
	
	_acc_skin_count=0.0;
	_acc_bg_count=0.0;
}

void SkinModel::Init(istream &s){
	int r_resolution, g_resolution;
	
	/*read the histogram resolutioin*/
	s>>r_resolution;
	s>>g_resolution;

	/*construct empty model*/
	Init(r_resolution, g_resolution);
	
	/*read skin data*/
	for(int i=0;i<COLOR_DIMENSIONS;i++){
		s>>_skin_mean[i];
		s>>_skin_var[i];
	}
	s>>_skin_prior;
	
	/*read background data*/
	for(int i=0;i<_hist_resolution[0];i++)
		for(int j=0;j<_hist_resolution[1];j++){
			s>>_bg[i][j];
		}
	
	/*update params for evaluation*/
	ProbabilityUpdateRetrained();
}

void SkinModel::Init(char const *filename){
	ifstream f(filename);
	
	if(!f.is_open())
		throw OpenFileException(filename);
	
	Init(f);
	
	f.close();
}

SkinModel::SkinModel(int r_resolution, int g_resolution){
	Init(r_resolution, g_resolution);
}

SkinModel::SkinModel(istream &s){
	Init(s);
}

SkinModel::SkinModel(char const *filename){
	Init(filename);
}


SkinModel::~SkinModel(){
	for(int i=0;i<_hist_resolution[0];i++){
		delete[] _p_diff[i];
		delete[] _bg[i];
		delete[] _acc_bg[i];
	}
	
	delete[] _p_diff;
	delete[] _bg;
	delete[] _acc_bg;
}

void SkinModel::AddSkin(int r, int g, int b, double weight){
	double value[COLOR_DIMENSIONS];

	if(TransformColor(r,g,b,value)){		
		/*add weightet value to the sum and squared sum*/
		for(int i=0;i<COLOR_DIMENSIONS;i++){
			_acc_skin_sum[i]+=value[i]*weight;
			_acc_skin_sqrsum[i]+=value[i]*value[i]*weight;
		}
		
		/*increase sum by the weight*/
		_acc_skin_count+=weight;
	}
}

void SkinModel::AddBackground(int r, int g, int b, double weight){
	double value[COLOR_DIMENSIONS];
	
	if(TransformColor(r,g,b,value)){		
		int bin[COLOR_DIMENSIONS];
	
		/*calculate bin indexes*/
		for(int i=0;i<COLOR_DIMENSIONS;i++)
			bin[i]=VALUE_TO_BIN(value[i],_hist_resolution[i]);
		
		/*add value to histograms*/
		_acc_bg[bin[0]][bin[1]]+=weight;
		_acc_bg_count+=weight;
	}
}

void SkinModel::RememberOnly(double final_count){
	/*skin*/
	if(_acc_bg_count>final_count){
		for(int i=0;i<COLOR_DIMENSIONS;i++){
			_acc_skin_sum[i]*=final_count/_acc_skin_count;
			_acc_skin_sqrsum[i]*=final_count/_acc_skin_count;
		}
		
		_acc_skin_count=final_count;
	}
	
	/*background*/
	if(_acc_bg_count>final_count){
		for(int i=0;i<_hist_resolution[0];i++)
				for(int j=0;j<_hist_resolution[1];j++)
					_acc_bg[i][j]*=final_count/_acc_bg_count;
					
		_acc_bg_count=final_count;
	}
}

void SkinModel::RememberRatio(double ratio){
	/*skin*/
	for(int i=0;i<COLOR_DIMENSIONS;i++){
		_acc_skin_sum[i]*=ratio;
		_acc_skin_sqrsum[i]*=ratio;
	}
	
	_acc_skin_count*=ratio;
	
	/*background*/
	for(int i=0;i<_hist_resolution[0];i++)
			for(int j=0;j<_hist_resolution[1];j++)
				_acc_bg[i][j]*=ratio;
				
	_acc_bg_count*=ratio;
}

void SkinModel::Retrain(){
	/*retrain skin*/
	if(_acc_skin_count>1.0)
		for(int i=0;i<COLOR_DIMENSIONS;i++){
			_skin_mean[i]=_acc_skin_sum[i]/_acc_skin_count;
			_skin_var[i]=_acc_skin_sqrsum[i]/_acc_skin_count-_skin_mean[i]*_skin_mean[i];
		}
	
	/*retrain background*/
	if(_acc_bg_count>1.0){	
		FOR_ALL_HIST_BINS{
			_bg[HIST_INDEX]=_acc_bg[HIST_INDEX]/_acc_bg_count;
		}
	}
		
	/*update parameers used for evaluation*/
	ProbabilityUpdateRetrained();
}

void SkinModel::Adapt(double tau){
	ProbabilityUpdateAdapted(tau);
}	

double SkinModel::Eval(int r, int g, int b) const{
	double value[COLOR_DIMENSIONS];
	
	if(TransformColor(r,g,b,value)){		
		int bin[COLOR_DIMENSIONS];
		
		/*calculate bin indexes*/
		for(int i=0;i<COLOR_DIMENSIONS;i++)
			bin[i]=VALUE_TO_BIN(value[i],_hist_resolution[i]);
			
		return _p_diff[bin[0]][bin[1]];
	}
	
	return 0.0;
}

void SkinModel::Eval(IplImage const *src, IplImage *dst){
	#if 0
	for(int i=0;i<_hist_resolution[0];i++)
		for(int j=0;j<_hist_resolution[1];j++){
			unsigned char *pixel=IMG_ELEM_NU(dst,i,j);
			
			int color=cvRound(_p_diff[i][j]*255);
			
			for(int i=0;i<dst->nChannels;i++)
				pixel[i]=color;
		}
	return;
	#endif
	
	for(int y=0;y<src->height;y++){
		unsigned char *src_pixel=((unsigned char*)(src->imageData + y*src->widthStep));
		unsigned char *dst_pixel=((unsigned char*)(dst->imageData + y*dst->widthStep));
			
		for(int x=0;x<src->width;x++){
			double value=Eval(src_pixel[R_OFFSET],src_pixel[G_OFFSET],src_pixel[B_OFFSET]);

			int color=cvRound(value*255);
			/*	
			if(value<0.5)
				color=0;
			else 
				color=255;				
			*/	
			for(int i=0;i<dst->nChannels;i++)
				dst_pixel[i]=color;
					
			/*next pixel*/
			src_pixel+=src->nChannels;
			dst_pixel+=dst->nChannels;
		}
	}
}
	
void SkinModel::Save(ostream &s) const{
	/*histogram size*/
	for(int i=0;i<COLOR_DIMENSIONS;i++)
		s<<_hist_resolution[i]<<" ";
		
	s<<endl<<endl;
	
	/*skin values*/
	for(int i=0;i<COLOR_DIMENSIONS;i++){
		s<<_skin_mean[i]<<" "<<_skin_var[i]<<endl;
	}
	s<<_skin_prior<<endl<<endl;
	
	/*background values*/
	for(int i=0;i<_hist_resolution[0];i++){
		for(int j=0;j<_hist_resolution[1];j++)
			s<<_bg[i][j]<<"\t";
			
		s<<endl;
	}
}
