/*!
  \file
  \brief Bayesian Filtering using stochastic sampling (Particle Filters)
  \author Vaclav Smidl.

  -----------------------------------
  BDM++ - C++ library for Bayesian Decision Making under Uncertainty

  Using IT++ for numerical operations
  -----------------------------------
*/

#ifndef PARTICLES_H
#define PARTICLES_H


#include "../estim/arx_ext.h"
#include "../stat/emix.h"

namespace bdm {
	
//! class used in PF
class MarginalizedParticle : public BM{
	protected:
	//! discrte particle
	dirac est_emp;
	//! internal Bayes Model
	shared_ptr<BM> bm; 
	//! pdf with for transitional par
	shared_ptr<pdf> par; // pdf for non-linear part
	//! link from this to bm
	shared_ptr<datalink_part> cond2bm;
	//! link from cond to par
	shared_ptr<datalink_part> cond2par;
	//! link from emp 2 par
	shared_ptr<datalink_part> emp2bm;
	//! link from emp 2 par
	shared_ptr<datalink_part> emp2par;
	
	//! custom posterior - product
	class eprod_2:public eprod_base {
		protected:
		MarginalizedParticle &mp;
		public:
		eprod_2(MarginalizedParticle &m):mp(m){}
		const epdf* factor(int i) const {return (i==0) ? &mp.bm->posterior() : &mp.est_emp;}
		const int no_factors() const {return 2;}
	} est;
	
	public:
		MarginalizedParticle():est(*this){};
		MarginalizedParticle(const MarginalizedParticle &m2):est(*this){
			bm = m2.bm->_copy();
			est_emp = m2.est_emp;
			par = m2.par;
			validate();
		};
		BM* _copy() const{return new MarginalizedParticle(*this);};
		void bayes(const vec &dt, const vec &cond){
			vec par_cond(par->dimensionc());
			cond2par->filldown(cond,par_cond); // copy ut
			emp2par->filldown(est_emp._point(),par_cond); // copy xt-1
			
			//sample new particle
			est_emp.set_point(par->samplecond(par_cond));
			//if (evalll)
			vec bm_cond(bm->dimensionc());
			cond2bm->filldown(cond, bm_cond);// set e.g. ut
			emp2bm->filldown(est_emp._point(), bm_cond);// set e.g. ut
			bm->bayes(dt, bm_cond);
			ll=bm->_ll();
		}
		const eprod_2& posterior() const {return est;}
	
		void set_prior(const epdf *pdf0){
			const eprod *ep=dynamic_cast<const eprod*>(pdf0);
			if (ep){ // full prior
				bdm_assert(ep->no_factors()==2,"Incompatible prod");
				bm->set_prior(ep->factor(0));
				est_emp.set_point(ep->factor(1)->sample());
			} else {
				// assume prior is only for emp;
				est_emp.set_point(pdf0->sample());
			}
		}

		/*! parse structure
		\code
		class = "BootstrapParticle";
		parameter_pdf = {class = 'epdf_offspring', ...};
		observation_pdf = {class = 'epdf_offspring',...};
		\endcode
		If rvs are set, then it checks for compatibility.
		*/
		void from_setting(const Setting &set){
			BM::from_setting ( set );
			par = UI::build<pdf> ( set, "parameter_pdf", UI::compulsory );
			bm = UI::build<BM> ( set, "bm", UI::compulsory );
		}
		
		void to_setting(const Setting &set){
			if (BM::log_level[logfull]){
			}
		}
		void validate(){
			if (est_emp.point.length()!=par->dimension())
				est_emp.set_point(zeros(par->dimension()));
			est.validate();
			
			yrv = bm->_yrv();
			dimy = bm->dimensiony();
			set_rv( concat(bm->_rv(), par->_rv()));
			set_dim( par->dimension()+bm->dimension());

			rvc = par->_rvc();
			rvc.add(bm->_rvc());
			rvc=rvc.subt(par->_rv());
			rvc=rvc.subt(par->_rv().copy_t(-1));
			rvc=rvc.subt(bm->_rv().copy_t(-1)); //
			
			cond2bm=new datalink_part;
			cond2par=new datalink_part;
			emp2bm  =new datalink_part;
			emp2par =new datalink_part;
			cond2bm->set_connection(bm->_rvc(), rvc);
			cond2par->set_connection(par->_rvc(), rvc);
			emp2bm->set_connection(bm->_rvc(), par->_rv());
			emp2par->set_connection(par->_rvc(), par->_rv().copy_t(-1));
			
			dimc = rvc._dsize();
		};
};
UIREGISTER(MarginalizedParticle);

//! class used in PF
class BootstrapParticle : public BM{
	dirac est;
	shared_ptr<pdf> par;
	shared_ptr<pdf> obs;
	shared_ptr<datalink_part> cond2par;
	shared_ptr<datalink_part> cond2obs;
	shared_ptr<datalink_part> xt2obs;
	shared_ptr<datalink_part> xtm2par;
	public:
		BM* _copy() const{return new BootstrapParticle(*this);};
		void bayes(const vec &dt, const vec &cond){
			vec par_cond(par->dimensionc());
			cond2par->filldown(cond,par_cond); // copy ut
			xtm2par->filldown(est._point(),par_cond); // copy xt-1
			
			//sample new particle
			est.set_point(par->samplecond(par_cond));
			//if (evalll)
			vec obs_cond(obs->dimensionc());
			cond2obs->filldown(cond, obs_cond);// set e.g. ut
			xt2obs->filldown(est._point(), obs_cond);// set e.g. ut
			ll=obs->evallogcond(dt,obs_cond);
		}
		const dirac& posterior() const {return est;}
		
		void set_prior(const epdf *pdf0){est.set_point(pdf0->sample());}
		
		/*! parse structure
		\code
		class = "BootstrapParticle";
		parameter_pdf = {class = 'epdf_offspring', ...};
		observation_pdf = {class = 'epdf_offspring',...};
		\endcode
		If rvs are set, then it checks for compatibility.
		*/
		void from_setting(const Setting &set){
			BM::from_setting ( set );
			par = UI::build<pdf> ( set, "parameter_pdf", UI::compulsory );
			obs = UI::build<pdf> ( set, "observation_pdf", UI::compulsory );
		}
		void validate(){
			yrv = obs->_rv();
			dimy = obs->dimension();
			set_rv( par->_rv());
			set_dim( par->dimension());
			
			rvc = par->_rvc().subt(par->_rv().copy_t(-1));
			rvc.add(obs->_rvc()); //
			
			cond2obs=new datalink_part;
			cond2par=new datalink_part;
			xt2obs  =new datalink_part;
			xtm2par =new datalink_part;
			cond2obs->set_connection(obs->_rvc(), rvc);
			cond2par->set_connection(par->_rvc(), rvc);
			xt2obs->set_connection(obs->_rvc(), _rv());
			xtm2par->set_connection(par->_rvc(), _rv().copy_t(-1));
			
			dimc = rvc._dsize();
		};
};
UIREGISTER(BootstrapParticle);


/*!
* \brief Trivial particle filter with proposal density equal to parameter evolution model.

Posterior density is represented by a weighted empirical density (\c eEmp ).
*/

class PF : public BM {
	//! \var log_level_enums weights
	//! all weightes will be logged

	//! \var log_level_enums menas
	//! means of particles will be logged
	LOG_LEVEL(PF,logweights,logmeans,logvars);
	
	class pf_mix: public emix_base{
		Array<BM*> &bms;
		public:
			pf_mix(vec &w0, Array<BM*> &bms0):emix_base(w0),bms(bms0){}
			const epdf* component(const int &i)const{return &(bms(i)->posterior());}
			int no_coms() const {return bms.length();}
	};
protected:
	//!number of particles;
	int n;
	//!posterior density
	pf_mix est;
	//! weights;
	vec w;
	//! particles
	Array<BM*> particles;
	//! internal structure storing loglikelihood of predictions
	vec lls;

	//! which resampling method will be used
	RESAMPLING_METHOD resmethod;
	//! resampling threshold; in this case its meaning is minimum ratio of active particles
	//! For example, for 0.5 resampling is performed when the numebr of active aprticles drops belo 50%.
	double res_threshold;

	//! \name Options
	//!@{
	//!@}

public:
	//! \name Constructors
	//!@{
	PF ( ) : est(w,particles) { };

	void set_parameters ( int n0, double res_th0 = 0.5, RESAMPLING_METHOD rm = SYSTEMATIC ) {
		n = n0;
		res_threshold = res_th0;
		resmethod = rm;
	};
	void set_model ( const BM *particle0, const epdf *prior) {
		if (n>0){
			particles.set_length(n);
			for (int i=0; i<n;i++){
				particles(i) = particle0->_copy();
				particles(i)->set_prior(prior);
			}
		}
		// set values for posterior
		est.set_rv ( particle0->posterior()._rv() );
	};
	void set_statistics ( const vec w0, const epdf &epdf0 ) {
		//est.set_statistics ( w0, epdf0 );
	};
/*	void set_statistics ( const eEmp &epdf0 ) {
		bdm_assert_debug ( epdf0._rv().equal ( par->_rv() ), "Incompatible input" );
		est = epdf0;
	};*/
	//!@}

	//! bayes compute weights of the
	virtual void bayes_weights();
	//! important part of particle filtering - decide if it is time to perform resampling
	virtual bool do_resampling() {
		double eff = 1.0 / ( w * w );
		return eff < ( res_threshold*n );
	}
	void bayes ( const vec &yt, const vec &cond );
	//!access function
	vec& _lls() {
		return lls;
	}
	//!access function
	RESAMPLING_METHOD _resmethod() const {
		return resmethod;
	}
	//! return correctly typed posterior (covariant return)
	const pf_mix& posterior() const {
		return est;
	}

	/*! configuration structure for basic PF
	\code
	parameter_pdf   = pdf_class;         // parameter evolution pdf
	observation_pdf = pdf_class;         // observation pdf
	prior           = epdf_class;         // prior probability density
	--- optional ---
	n               = 10;                 // number of particles
	resmethod       = 'systematic', or 'multinomial', or 'stratified'
										  // resampling method
	res_threshold   = 0.5;                // resample when active particles drop below 50%
	\endcode
	*/
	void from_setting ( const Setting &set ) {
		BM::from_setting ( set );
		UI::get ( log_level, set, "log_level", UI::optional );
		
		shared_ptr<BM> bm0 = UI::build<BM>(set, "particle",UI::compulsory);
		
		shared_ptr<epdf> pri = UI::build<epdf> ( set, "prior", UI::compulsory );
		n =0;
		UI::get(n,set,"n",UI::optional);;
		if (n>0){
			particles.set_length(n);
			for(int i=0;i<n;i++){particles(i)=bm0->_copy();}
			w = ones(n)/n;
		}
		set_prior(pri.get());
		// set resampling method
		resmethod_from_set ( set );
		//set drv

		rvc = bm0->_rvc();
		dimc = bm0->dimensionc();
		BM::set_rv(bm0->_rv());
		yrv=bm0->_yrv();
		dimy = bm0->dimensiony();
	}
	
	void log_register ( bdm::logger& L, const string& prefix ){
		BM::log_register(L,prefix);
		if (log_level[logweights]){
			L.add_vector( log_level, logweights, RV ( particles.length()), prefix); 
		}
		if (log_level[logmeans]){
			for (int i=0; i<particles.length(); i++){
				L.add_vector( log_level, logmeans, RV ( particles(i)->dimension() ), prefix , i);
			}
		}
		if (log_level[logvars]){
			for (int i=0; i<particles.length(); i++){
				L.add_vector( log_level, logvars, RV ( particles(i)->dimension() ), prefix , i);
			}
		}
	};
	void log_write ( ) const {
		BM::log_write();
		if (log_level[logweights]){
			log_level.store( logweights, w); 
		}
		if (log_level[logmeans]){
			for (int i=0; i<particles.length(); i++){
				log_level.store( logmeans, particles(i)->posterior().mean(), i);
			}
		}
		if (log_level[logvars]){
			for (int i=0; i<particles.length(); i++){
				log_level.store( logvars, particles(i)->posterior().variance(), i);
			}
		}
		
	}
	
	void set_prior(const epdf *pri){
		const emix_base *emi=dynamic_cast<const emix_base*>(pri);
		if (emi) {
			bdm_assert(particles.length()>0, "initial particle is not assigned");
			n = emi->_w().length();
			int old_n = particles.length();
			if (n!=old_n){
				particles.set_length(n,true);
			} 
			for(int i=old_n;i<n;i++){particles(i)=particles(0)->_copy();}
			
			for (int i =0; i<n; i++){
				particles(i)->set_prior(emi->_com(i));
			}
		} else {
			// try to find "n"
			bdm_assert(n>0, "Field 'n' must be filled when prior is not of type emix");
			for (int i =0; i<n; i++){
				particles(i)->set_prior(pri);
			}
			
		}
	}
	//! auxiliary function reading parameter 'resmethod' from configuration file
	void resmethod_from_set ( const Setting &set ) {
		string resmeth;
		if ( UI::get ( resmeth, set, "resmethod", UI::optional ) ) {
			if ( resmeth == "systematic" ) {
				resmethod = SYSTEMATIC;
			} else  {
				if ( resmeth == "multinomial" ) {
					resmethod = MULTINOMIAL;
				} else {
					if ( resmeth == "stratified" ) {
						resmethod = STRATIFIED;
					} else {
						bdm_error ( "Unknown resampling method" );
					}
				}
			}
		} else {
			resmethod = SYSTEMATIC;
		};
		if ( !UI::get ( res_threshold, set, "res_threshold", UI::optional ) ) {
			res_threshold = 0.9;
		}
		//validate();
	}

	void validate() {
		BM::validate();
		est.validate();
		bdm_assert ( n>0, "empty particle pool" );
		n = w.length();
		lls = zeros ( n );

		if ( particles(0)->_rv()._dsize() > 0 ) {
			bdm_assert (  particles(0)->_rv()._dsize() == est.dimension(), "Mismatch of RV and dimension of posterior" );
		}
	}
	//! resample posterior density (from outside - see MPF)
	void resample ( ) {
		ivec ind = zeros_i ( n );
		bdm::resample(w,ind,resmethod);
		// copy the internals according to ind
		for (int i = 0; i < n; i++ ) {
			if ( ind ( i ) != i ) {
				particles( i ) = particles( ind ( i ) )->_copy();
			}
			w ( i ) = 1.0 / n;
		}
	}
	//! access function
	Array<BM*>& _particles() {
		return particles;
	}

};
UIREGISTER ( PF );

/*!
\brief Marginalized Particle filter

A composition of particle filter with exact (or almost exact) bayesian models (BMs).
The Bayesian models provide marginalized predictive density. Internaly this is achieved by virtual class MPFpdf.
*/

// class MPF : public BM  {
// 	//! Introduces new option 
// 	//! \li means - meaning TODO
// 	LOG_LEVEL(MPF,means);
// protected:
// 	//! particle filter on non-linear variable
// 	shared_ptr<PF> pf;
// 	//! Array of Bayesian models
// 	Array<BM*> BMs;
// 
// 	//! internal class for pdf providing composition of eEmp with external components
// 
// 	class mpfepdf : public epdf  {
// 		//! pointer to particle filter
// 		shared_ptr<PF> &pf;
// 		//! pointer to Array of BMs
// 		Array<BM*> &BMs;
// 	public:
// 		//! constructor
// 		mpfepdf ( shared_ptr<PF> &pf0, Array<BM*> &BMs0 ) : epdf(), pf ( pf0 ), BMs ( BMs0 ) { };
// 		//! a variant of set parameters - this time, parameters are read from BMs and pf
// 		void read_parameters() {
// 			rv = concat ( pf->posterior()._rv(), BMs ( 0 )->posterior()._rv() );
// 			dim = pf->posterior().dimension() + BMs ( 0 )->posterior().dimension();
// 			bdm_assert_debug ( dim == rv._dsize(), "Wrong name " );
// 		}
// 		vec mean() const;
// 
// 		vec variance() const;
// 
// 		void qbounds ( vec &lb, vec &ub, double perc = 0.95 ) const;
// 
// 		vec sample() const NOT_IMPLEMENTED(0);
// 
// 		double evallog ( const vec &val ) const NOT_IMPLEMENTED(0);		
// 	};
// 
// 	//! Density joining PF.est with conditional parts
// 	mpfepdf jest;
// 
// 	//! datalink from global yt and cond (Up) to BMs yt and cond (Down)
// 	datalink_m2m this2bm;
// 	//! datalink from global yt and cond (Up) to PFs yt and cond (Down)
// 	datalink_m2m this2pf;
// 	//!datalink from PF part to BM
// 	datalink_part pf2bm;
// 
// public:
// 	//! Default constructor.
// 	MPF () :  jest ( pf, BMs ) {};
// 	//! set all parameters at once
// 	void set_pf ( shared_ptr<pdf> par0, int n0, RESAMPLING_METHOD rm = SYSTEMATIC ) {
// 		if (!pf) pf=new PF;
// 		pf->set_model ( par0, par0 ); // <=== nasty!!!
// 		pf->set_parameters ( n0, rm );
// 		pf->set_rv(par0->_rv());
// 		BMs.set_length ( n0 );
// 	}
// 	//! set a prototype of BM, copy it to as many times as there is particles in pf
// 	void set_BM ( const BM &BMcond0 ) {
// 
// 		int n = pf->__w().length();
// 		BMs.set_length ( n );
// 		// copy
// 		//BMcond0 .condition ( pf->posterior()._sample ( 0 ) );
// 		for ( int i = 0; i < n; i++ ) {
// 			BMs ( i ) = (BM*) BMcond0._copy();
// 		}
// 	};
// 
// 	void bayes ( const vec &yt, const vec &cond );
// 
// 	const epdf& posterior() const {
// 		return jest;
// 	}
// 
// 	//!Access function
// 	const BM* _BM ( int i ) {
// 		return BMs ( i );
// 	}
// 	PF& _pf() {return *pf;}
// 
// 
// 	virtual double logpred ( const vec &yt ) const NOT_IMPLEMENTED(0); 
// 		
// 	virtual epdf* epredictor() const NOT_IMPLEMENTED(NULL); 
// 	
// 	virtual pdf* predictor() const NOT_IMPLEMENTED(NULL); 
// 
// 
// 	/*! configuration structure for basic PF
// 	\code
// 	BM              = BM_class;           // Bayesian filtr for analytical part of the model
// 	parameter_pdf   = pdf_class;         // transitional pdf for non-parametric part of the model
// 	prior           = epdf_class;         // prior probability density
// 	--- optional ---
// 	n               = 10;                 // number of particles
// 	resmethod       = 'systematic', or 'multinomial', or 'stratified'
// 										  // resampling method
// 	\endcode
// 	*/
// 	void from_setting ( const Setting &set ) {
// 		BM::from_setting( set );
// 
// 		shared_ptr<pdf> par = UI::build<pdf> ( set, "parameter_pdf", UI::compulsory );
// 
// 		pf = new PF;
// 		// prior must be set before BM
// 		pf->prior_from_set ( set );
// 		pf->resmethod_from_set ( set );
// 		pf->set_model ( par, par ); // too hackish!
// 
// 		shared_ptr<BM> BM0 = UI::build<BM> ( set, "BM", UI::compulsory );
// 		set_BM ( *BM0 );
// 
// 		//set drv
// 		//??set_yrv(concat(BM0->_yrv(),u) );
// 		set_yrv ( BM0->_yrv() );
// 		rvc = BM0->_rvc().subt ( par->_rv() );
// 		//find potential input - what remains in rvc when we subtract rv
// 		RV u = par->_rvc().subt ( par->_rv().copy_t ( -1 ) );
// 		rvc.add ( u );
// 		dimc = rvc._dsize();
// 		validate();
// 	}
// 
// 	void validate() {
// 		BM::validate();
// 		try {
// 			pf->validate();
// 		} catch ( std::exception ) {
// 			throw UIException ( "Error in PF part of MPF:" );
// 		}
// 		jest.read_parameters();
// 		this2bm.set_connection ( BMs ( 0 )->_yrv(), BMs ( 0 )->_rvc(), yrv, rvc );
// 		this2pf.set_connection ( pf->_yrv(), pf->_rvc(), yrv, rvc );
// 		pf2bm.set_connection ( BMs ( 0 )->_rvc(), pf->posterior()._rv() );
// 	}
// };
// UIREGISTER ( MPF );

/*! ARXg for estimation of state-space variances
*/
// class MPF_ARXg :public BM{
// 	protected:
// 	shared_ptr<PF> pf;
// 	//! pointer to Array of BMs
// 	Array<ARX*> BMso;
// 	//! pointer to Array of BMs
// 	Array<ARX*> BMsp;
// 	//!parameter evolution
// 	shared_ptr<fnc> g;
// 	//!observation function
// 	shared_ptr<fnc> h;
// 	
// 	public:
// 		void bayes(const vec &yt, const vec &cond );
// 		void from_setting(const Setting &set) ;
// 		void validate() {
// 			bdm_assert(g->dimensionc()==g->dimension(),"not supported yet");
// 			bdm_assert(h->dimensionc()==g->dimension(),"not supported yet");			
// 		}
// 
// 		double logpred(const vec &cond) const NOT_IMPLEMENTED(0.0);
// 		epdf* epredictor() const NOT_IMPLEMENTED(NULL);
// 		pdf* predictor() const NOT_IMPLEMENTED(NULL);
// 		const epdf& posterior() const {return pf->posterior();};
// 		
// 		void log_register( logger &L, const string &prefix ){
// 			BM::log_register(L,prefix);
// 			registered_logger->ids.set_size ( 3 );
// 			registered_logger->ids(1)= L.add_vector(RV("Q",dimension()*dimension()), prefix+L.prefix_sep()+"Q");
// 			registered_logger->ids(2)= L.add_vector(RV("R",dimensiony()*dimensiony()), prefix+L.prefix_sep()+"R");
// 			
// 		};
// 		void log_write() const {
// 			BM::log_write();
// 			mat mQ=zeros(dimension(),dimension());
// 			mat pom=zeros(dimension(),dimension());
// 			mat mR=zeros(dimensiony(),dimensiony());
// 			mat pom2=zeros(dimensiony(),dimensiony());
// 			mat dum;
// 			const vec w=pf->posterior()._w();
// 			for (int i=0; i<w.length(); i++){
// 				BMsp(i)->posterior().mean_mat(dum,pom);
// 				mQ += w(i) * pom;
// 				BMso(i)->posterior().mean_mat(dum,pom2);
// 				mR += w(i) * pom2;
// 				
// 			}
// 			registered_logger->L.log_vector ( registered_logger->ids ( 1 ), cvectorize(mQ) );
// 			registered_logger->L.log_vector ( registered_logger->ids ( 2 ), cvectorize(mR) );
// 			
// 		}
// };
// UIREGISTER(MPF_ARXg);


}
#endif // KF_H


