/*!
  \file
  \brief Bayesian Filtering for generalized autoregressive (ARX) model
  \author Vaclav Smidl.

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

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

#ifndef AR_H
#define AR_H

#include "../math/functions.h"
#include "../stat/exp_family.h"
#include "../base/user_info.h"
//#include "../estim/kalman.h"
#include "arx_straux.h"

namespace bdm {

/*!
* \brief Linear Autoregressive model with Gaussian noise

Regression of the following kind:
\f[
y_t = \theta_1 \psi_1 + \theta_2 + \psi_2 +\ldots + \theta_n \psi_n + r e_t
\f]
where unknown parameters \c rv are \f$[\theta r]\f$, regression vector \f$\psi=\psi(y_{1:t},u_{1:t})\f$ is a known function of past outputs and exogeneous variables \f$u_t\f$. Distrubances \f$e_t\f$ are supposed to be normally distributed:
\f[
e_t \sim \mathcal{N}(0,1).
\f]

See \ref tut_arx for mathematical treatment.

The easiest way how to use the class is:
\include arx_simple.cpp

	\todo sort out constant terms - bayes should accept vec without additional 1s
*/
class ARX: public BMEF {
protected:
	//! switch if constant is modelled or not
	bool have_constant;
	//! vector of dyadic update
	vec dyad;
	//! posterior density
	egiw est;
	//! Alternative estimate of parameters, used in stabilized forgetting, see [Kulhavy]
	egiw alter_est;
public:
	//! \name Constructors
	//!@{
	ARX ( const double frg0 = 1.0 ) : BMEF ( frg0 ),  have_constant ( true ), dyad(), est(), alter_est() {};
	ARX ( const ARX &A0 ) : BMEF ( A0 ),  have_constant ( A0.have_constant ), dyad ( A0.dyad ), est ( A0.est ), alter_est ( A0.alter_est ) { };

	ARX* _copy() const;

	void set_frg ( double frg0 ) {
		frg = frg0;
	}
	void set_constant ( bool const0 ) {
		have_constant = const0;
	}
	void set_statistics ( int dimy0, const ldmat V0, double nu0 = -1.0 ) {
		est.set_parameters ( dimy0, V0, nu0 );
		est.validate();
		last_lognc = est.lognc();
		dimy = dimy0;
	}
	//!@}

	//! Set sufficient statistics
	void set_statistics ( const BMEF* BM0 );

	//!\name Mathematical operations
	//!@{

	//! Weighted Bayes \f$ dt = [y_t psi_t] \f$.
	void bayes_weighted ( const vec &yt, const vec &cond = empty_vec, const double w = 1.0 );
	void bayes ( const vec &yt, const vec &cond = empty_vec ) {
		bayes_weighted ( yt, cond, 1.0 );
	};
	double logpred ( const vec &yt ) const;
	void flatten ( const BMEF* B );
	//! Conditioned version of the predictor
	enorm<ldmat>* epredictor ( const vec &rgr ) const;
	//! Predictor for empty regressor
	enorm<ldmat>* epredictor() const;
	//! conditional version of the predictor
	template<class sq_T>
	shared_ptr<mlnorm<sq_T> > ml_predictor() const;
	//! fast version of predicto
	template<class sq_T>
	void ml_predictor_update ( mlnorm<sq_T> &pred ) const;
	mlstudent* predictor_student() const;
	//! Brute force structure estimation.\return indeces of accepted regressors.
	ivec structure_est ( egiw Eg0 );
	//! Smarter structure estimation by Ludvik Tesar.\return indeces of accepted regressors.
	ivec structure_est_LT ( egiw Eg0 );
	//!@}

	//!\name Access attributes
	//!@{
	//! return correctly typed posterior (covariant return)
	const egiw& posterior() const {
		return est;
	}
	//!@}

	/*! UI for ARX estimator

	\code
	class = 'ARX';
	rv    = RV({names_of_dt} )                 // description of output variables
	rgr   = RV({names_of_regressors}, [-1,-2]} // description of regressor variables
	constant = 1;                              // 0/1 switch if the constant term is modelled or not

	--- optional ---
	prior = {class='egiw',...};                // Prior density, when given default is used instead
	alternative = {class='egiw',...};          // Alternative density in stabilized estimation, when not given prior is used

	frg = 1.0;                                 // forgetting, default frg=1.0

	rv_param   = RV({names_of_parameters}}     // description of parametetr names
											   // default: ["theta_i" and "r_i"]
	\endcode
	*/
	void from_setting ( const Setting &set );

	void validate() {
		BMEF::validate();	

		//if dimc not set set it from V
		if ( dimc == 0 ) {
			dimc = posterior()._V().rows() - dimy - int ( have_constant == true );
		}

		if ( have_constant ) {
			dyad = ones ( dimy + dimc + 1 );
		} else {
			dyad = zeros ( dimy + dimc );
		}

	}
	//! function sets prior and alternative density
	void set_prior ( const RV &drv, egiw &prior ) {
		//TODO check ranges in RV and build prior
	};
	//! build default prior and alternative when all values are set
	void set_prior_default ( egiw &prior ) {
		//assume
		vec dV0 ( prior._V().rows() );
		dV0.set_subvector ( 0, prior._dimx() - 1, 1.0 );
		if (dV0.length()>prior._dimx())
			dV0.set_subvector ( prior._dimx(), dV0.length() - 1, 1e-5 );
		
		prior.set_parameters ( prior._dimx(), ldmat ( dV0 ) );
		prior.validate();
	}

	void to_setting ( Setting &set ) const
	{			
		BMEF::to_setting( set ); // takes care of rv, yrv, rvc
		int constant = have_constant ? 1 : 0;
		UI::save(constant, set, "constant");
		UI::save(&est, set, "prior");
		UI::save(&alter_est, set, "alternative");
		
		
	} 
};
UIREGISTER ( ARX );
SHAREDPTR ( ARX );

/*! ARX model conditined by knowledge of the forgetting factor
\f[ f(\theta| d_1 \ldots d_t , \phi_t) \f]

The symbol \f$ \phi \f$ is assumed to be the last of the conditioning variables.
*/
class ARXfrg : public ARX {
public:
	ARXfrg() : ARX() {};
	//! copy constructor
	ARXfrg ( const ARXfrg &A0 ) : ARX ( A0 ) {};
	virtual ARXfrg* _copy() const {
		ARXfrg *A = new ARXfrg ( *this );
		return A;
	}

	void bayes ( const vec &val, const vec &cond ) {
		frg = cond ( dimc - 1 ); //  last in cond is phi
		ARX::bayes ( val, cond );
	}
	void validate() {
		ARX::validate();
		rvc.add ( RV ( "{phi }", vec_1 ( 1 ) ) );
		dimc += 1;
	}
};
UIREGISTER ( ARXfrg );


////////////////////
template<class sq_T>
shared_ptr< mlnorm<sq_T> > ARX::ml_predictor ( ) const {
	shared_ptr< mlnorm<sq_T> > tmp = new mlnorm<sq_T> ( );
	tmp->set_rv ( yrv );
	tmp->set_rvc ( _rvc() );

	ml_predictor_update ( *tmp );
	tmp->validate();
	return tmp;
}

template<class sq_T>
void ARX::ml_predictor_update ( mlnorm<sq_T> &pred ) const {
	mat mu ( dimy, posterior()._V().rows() - dimy );
	mat R ( dimy, dimy );

	est.mean_mat ( mu, R ); //mu =
	mu = mu.T();
	//correction for student-t  -- TODO check if correct!!

	if ( have_constant ) { // constant term
		//Assume the constant term is the last one:
		pred.set_parameters ( mu.get_cols ( 0, mu.cols() - 2 ), mu.get_col ( mu.cols() - 1 ), sq_T ( R ) );
	} else {
		pred.set_parameters ( mu, zeros ( dimy ), sq_T ( R ) );
	}
}

};
#endif // AR_H


