/*!
  \file
  \brief Probability distributions for Mixtures of pdfs
  \author Vaclav Smidl.

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

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

#ifndef EMIX_H
#define EMIX_H

#define LOG2  0.69314718055995

#include "../shared_ptr.h"
#include "exp_family.h"

namespace bdm {

//this comes first because it is used inside emix!

/*! \brief Class representing ratio of two densities
which arise e.g. by applying the Bayes rule.
It represents density in the form:
\f[
f(rv|rvc) = \frac{f(rv,rvc)}{f(rvc)}
\f]
where \f$ f(rvc) = \int f(rv,rvc) d\ rv \f$.

In particular this type of arise by conditioning of a mixture model.

At present the only supported operation is evallogcond().
 */
class mratio: public pdf {
protected:
	//! Nominator in the form of pdf
	const epdf* nom;

	//!Denominator in the form of epdf
	shared_ptr<epdf> den;

	//!flag for destructor
	bool destroynom;
	//!datalink between conditional and nom
	datalink_m2e dl;
public:
	//!Default constructor. By default, the given epdf is not copied!
	//! It is assumed that this function will be used only temporarily.
	mratio ( const epdf* nom0, const RV &rv, bool copy = false ) : pdf ( ), dl ( ) {
		// adjust rv and rvc

		set_rv ( rv ); 
		dim = rv._dsize();

		rvc = nom0->_rv().subt ( rv );
		dimc = rvc._dsize();

		//prepare data structures
		if ( copy ) {
			bdm_error ( "todo" );
			// destroynom = true;
		} else {
			nom = nom0;
			destroynom = false;
		}
		bdm_assert_debug ( rvc.length() > 0, "Makes no sense to use this object!" );

		// build denominator
		den = nom->marginal ( rvc );
		dl.set_connection ( rv, rvc, nom0->_rv() );
	};

	double evallogcond ( const vec &val, const vec &cond ) {
		double tmp;
		vec nom_val ( dimension() + dimc );
		dl.pushup_cond ( nom_val, val, cond );
		tmp = exp ( nom->evallog ( nom_val ) - den->evallog ( cond ) );
		return tmp;
	}

	//! Returns a sample from the density conditioned on \c cond, \f$x \sim epdf(rv|cond)\f$. \param cond is numeric value of \c rv
	virtual vec samplecond ( const vec &cond ) NOT_IMPLEMENTED(0);

	//! Object takes ownership of nom and will destroy it
	void ownnom() {
		destroynom = true;
	}
	//! Default destructor
	~mratio() {
		if ( destroynom ) {
			delete nom;
		}
	}


private:
	// not implemented
	mratio ( const mratio & );
	mratio &operator= ( const mratio & );
};

class emix; //forward

/*! Base class (Interface) for mixtures 
*/
class emix_base : public epdf {
	protected:
	//! reference to vector of weights
	vec &w;
	//! function returning ith component
	virtual const epdf * component(const int &i) const=0;
	
	virtual int no_coms() const=0;
	
	public:
		
		emix_base(vec &w0): w(w0){}
		
		void validate ();
		
		vec sample() const;
		
		vec mean() const;
		
		vec variance() const;
		
		double evallog ( const vec &val ) const;
		
		vec evallog_mat ( const mat &Val ) const;
		
		//! Auxiliary function that returns pdflog for each component
		mat evallog_coms ( const mat &Val ) const;
		
		shared_ptr<epdf> marginal ( const RV &rv ) const;
		//! Update already existing marginal density  \c target
		void marginal ( const RV &rv, emix &target ) const;
		shared_ptr<pdf> condition ( const RV &rv ) const;
		
		//Access methods	
		//! returns a reference to the internal weights. Use with Care!
		vec& _w() {
			return w;
		}
		
		const vec& _w() const {
			return w;
		}
		//!access
		const epdf* _com(int i) const {return component(i);}
		
};

/*!
* \brief Mixture of epdfs

Density function:
\f[
f(x) = \sum_{i=1}^{n} w_{i} f_i(x), \quad \sum_{i=1}^n w_i = 1.
\f]
where \f$f_i(x)\f$ is any density on random variable \f$x\f$, called \a component,

*/
class emix : public emix_base {
protected:
	//! weights of the components
	vec weights;

	//! Component (epdfs)
	Array<shared_ptr<epdf> > Coms;

public:
	//! Default constructor
	emix ( ) : emix_base ( weights) { }
	
	const epdf* component(const int &i) const {return Coms(i).get();}
	void validate();
	

	int no_coms() const {return Coms.length(); }

	//! Load from structure with elements:
	//!  \code
	//! { class='emix';
	//!   pdfs = (..., ...);     // list of pdfs in the mixture
	//!   weights = ( 0.5, 0.5 ); // weights of pdfs in the mixture
	//! }
	//! \endcode
	//!@}
	void from_setting ( const Setting &set );
	void to_setting  (Setting  &set) const;

	void set_rv ( const RV &rv ) {
		epdf::set_rv ( rv );
		for ( int i = 0; i < no_coms(); i++ ) {
			Coms( i )->set_rv ( rv );
		}
	}
	
	Array<shared_ptr<epdf> >& _Coms ( ) {
			return Coms;
		}
};
SHAREDPTR ( emix );
UIREGISTER ( emix );


/*! \brief Chain rule decomposition of epdf

Probability density in the form of Chain-rule decomposition:
\[
f(x_1,x_2,x_3) = f(x_1|x_2,x_3)f(x_2,x_3)f(x_3)
\]
Note that
*/
class mprod: public pdf {
private:
	Array<shared_ptr<pdf> > pdfs;

	//! Data link for each pdfs
	Array<shared_ptr<datalink_m2m> > dls;

public:
	//! \brief Default constructor
	mprod() { }

	/*!\brief Constructor from list of mFacs
	*/
	mprod ( const Array<shared_ptr<pdf> > &mFacs ) {
		set_elements ( mFacs );
	}
	//! Set internal \c pdfs from given values
	void set_elements ( const Array<shared_ptr<pdf> > &mFacs );

	double evallogcond ( const vec &val, const vec &cond );

	vec evallogcond_mat ( const mat &Dt, const vec &cond );

	vec evallogcond_mat ( const Array<vec> &Dt, const vec &cond );

	//TODO smarter...
	vec samplecond ( const vec &cond ) {
		//! Ugly hack to help to discover if mpfs are not in proper order. Correct solution = check that explicitely.
		vec smp = std::numeric_limits<double>::infinity() * ones ( dimension() );
		vec smpi;
		// Hard assumption here!!! We are going backwards, to assure that samples that are needed from smp are already generated!
		for ( int i = ( pdfs.length() - 1 ); i >= 0; i-- ) {
			// generate contribution of this pdf
			smpi = pdfs ( i )->samplecond ( dls ( i )->get_cond ( smp , cond ) );
			// copy contribution of this pdf into smp
			dls ( i )->pushup ( smp, smpi );
		}
		return smp;
	}

	//! Load from structure with elements:
	//!  \code
	//! { class='mprod';
	//!   pdfs = (..., ...);     // list of pdfs in the order of chain rule
	//! }
	//! \endcode
	//!@}
	void from_setting ( const Setting &set ) ;
	void 	to_setting  (Setting  &set) const;


};
UIREGISTER ( mprod );
SHAREDPTR ( mprod );


//! Product of independent epdfs. For dependent pdfs, use mprod.
class eprod_base: public epdf {
protected:
	//! Array of indices
	Array<datalink*> dls;
	//! interface for a factor
public:
	virtual const epdf* factor(int i) const NOT_IMPLEMENTED(NULL); 
	//!number of factors
	virtual const int no_factors() const NOT_IMPLEMENTED(0);
	//! Default constructor
	eprod_base () :  dls (0) {};
	//! Set internal
	vec mean() const;

	vec variance() const;

	vec sample() const;

	double evallog ( const vec &val ) const;

	//!Destructor
	~eprod_base() {
		for ( int i = 0; i < dls.length(); i++ ) {
			delete dls ( i );
		}
	}
	void validate() {
		epdf::validate();
		dls.set_length ( no_factors() );
		
		bool independent = true;
		dim = 0;
		for ( int i = 0; i < no_factors(); i++ ) {
			independent = rv.add ( factor ( i )->_rv() );
			dim += factor ( i )->dimension();
			bdm_assert_debug ( independent, "eprod:: given components are not independent." );
		};
		
		//
		int cumdim = 0;
		int dimi = 0;
		int i;
		for ( i = 0; i < no_factors(); i++ ) {
			dls ( i ) = new datalink;
			if ( isnamed() ) { // rvs are complete
				dls ( i )->set_connection ( factor ( i )->_rv() , rv );
			} else { //rvs are not reliable
				dimi = factor ( i )->dimension();
				dls ( i )->set_connection ( dimi, dim, linspace ( cumdim, cumdim + dimi - 1 ) );
				cumdim += dimi;
			}
		}
	
	}
};

class eprod: public eprod_base{
	protected:
		Array<shared_ptr<epdf> > factors;
	public:
		const epdf* factor(int i) const {return factors(i).get();}
		const int no_factors() const {return factors.length();}
	void set_parameters ( const Array<shared_ptr<epdf> > &epdfs0) {
		factors = epdfs0;
	}
	void from_setting(const Setting &set){
		UI::get(factors,set,"pdfs",UI::compulsory);
	}
};
UIREGISTER(eprod);

//! similar to eprod but used only internally -- factors are external pointers
class eprod_internal: public eprod_base{
	protected:
		Array<epdf* > factors;
		const epdf* factor(int i) const {return factors(i);}
		const int no_factors() const {return factors.length();}
	public:
		void set_parameters ( const Array<epdf *> &epdfs0) {
			factors = epdfs0;
		}
};

/*! \brief Mixture of pdfs with constant weights, all pdfs are of equal RV and RVC

*/
class mmix : public pdf {
protected:
	//! Component (pdfs)
	Array<shared_ptr<pdf> > Coms;
	//!weights of the components
	vec w;
public:
	//!Default constructor
	mmix() : Coms ( 0 ) { };

	double evallogcond ( const vec &dt, const vec &cond ) {
		double ll = 0.0;
		for ( int i = 0; i < Coms.length(); i++ ) {
			ll += Coms ( i )->evallogcond ( dt, cond );
		}
		return ll;
	}

	vec samplecond ( const vec &cond );

	//! Load from structure with elements:
	//!  \code
	//! { class='mmix';
	//!   pdfs = (..., ...);     // list of pdfs in the mixture
	//!   weights = ( 0.5, 0.5 ); // weights of pdfs in the mixture
	//! }
	//! \endcode
	//!@}
	void from_setting ( const Setting &set );
	void 	to_setting  (Setting  &set) const;
	virtual void validate();
};
SHAREDPTR ( mmix );
UIREGISTER ( mmix );


//! Base class for all BM running as parallel update of internal BMs

class ProdBMBase : public BM {
	protected :
		Array<vec_from_vec> bm_yt;
		Array<vec_from_2vec> bm_cond;
		class eprod_bm : public eprod_base {
			ProdBMBase & pb;
			public :
			eprod_bm(ProdBMBase &pb0): pb(pb0){}
			const epdf* factor(int i ) const {return &(pb.bm(i)->posterior());}
			const int no_factors() const {return pb.no_bms();}
		} est;
	public:
		ProdBMBase():est(*this){}
		virtual BM* bm(int i) NOT_IMPLEMENTED(NULL);
		virtual int no_bms() const {return 0;}
		const epdf& posterior() const {return est;}
		void set_prior(const epdf *pri){
			const eprod_base* ep=dynamic_cast<const eprod_base*>(pri);
			if (ep){
				bdm_assert(ep->no_factors()!=no_bms() , "Given prior has "+ num2str(ep->no_factors()) + " while this ProdBM has "+
					num2str(no_bms()) + "BMs");
				for (int i=0; i<no_bms(); i++){
					bm(i)->set_prior(ep->factor(i));
				}
			}
		}
		
		void validate() {
			est.validate();
			BM::validate();
			// set links
			bm_yt.set_length(no_bms());
			bm_cond.set_length(no_bms());
			
			// 
			
			for (int i=0; i<no_bms(); i++){
				yrv.add(bm(i)->_yrv());
				rvc.add(bm(i)->_rvc());
			}
			rvc=rvc.subt(yrv);
			
			dimy = yrv._dsize();
			dimc = rvc._dsize();
			
			for (int i=0; i<no_bms(); i++){
				bm_yt(i).connect(bm(i)->_yrv(), yrv);
				bm_cond(i).connect(bm(i)->_rvc(), yrv, rvc);
			}
		}
		void bayes(const vec &dt, const vec &cond){
			ll=0;
			for(int i=0;i<no_bms(); i++){
				bm_yt(i).update(dt);
				bm_cond(i).update(dt,cond);
				bm(i)->bayes(bm_yt(i), bm_cond(i));
			}
		}
		
};

class ProdBM: public ProdBMBase{
	protected:
		Array<shared_ptr<BM> > BMs;
	public:
		virtual BM* bm(int i) {return BMs(i).get();}
		virtual int no_bms() const {return BMs.length();}
		void from_setting(const Setting &set){
			BM::from_setting(set);
			UI::get(BMs,set,"BMs");
		}
		void to_setting(Setting &set) const{
			BM::to_setting(set);
			UI::save(BMs,set,"BMs");
		}
		
};
UIREGISTER(ProdBM);

//! \brief class for on-line model selection 
class ModelComparator: public ProdBM{
	protected:
		multiBM weights;
	public:
		void bayes(const vec &yt, const vec &cond){
			vec w_nn(no_bms());
			for (int i=0; i<no_bms(); i++){
				bm(i)->bayes(yt,cond);
				w_nn(i) += bm(i)->_ll();
			}
			vec w=exp(w_nn-max(w_nn));
			weights.bayes(w/sum(w));
		}
		void validate(){
			ProdBM::validate();
			weights.validate();
		}
		void from_setting(const Setting& set){
			ProdBM::from_setting(set);
			UI::get(weights.frg, set, "frg",UI::optional);
		}
		void to_setting(Setting& set) const{
			ProdBM::to_setting(set);
			UI::save(weights.frg, set, "frg");
		}
};

}
#endif //MX_H
