#ifndef ARX_CTRL_H
#define ARX_CTRL_H

#include <design/ctrlbase.h>
using namespace bdm;

/*! exact controller for system \f$ y_{t+1} = y_{t} + b u_{t} + \sigma e_t \f$

From \f$ E\{y_{ref}\} = E\{ y_{t} + b u_{t} \} \f$, the optimal control action is:
\f[ u_t = (y_{ref} - y_{t})/b \f]

*/ 

class exact_ctrl: public Controller {
	protected:
		//! target value of y
		double yr;
		//! parameter b (or its estimate)
		double b;
	public:
		void set_b (const double b0){b=b0;}
		//! data in ctrlaction are only past outputs yt
		vec ctrlaction(const vec &yt){
			vec ut(1);
			if (abs(b)<1e-2) {
				ut= 0.1*randn();
			} else {
				ut = (yr-yt)/b;
			}
			return ut;
		}
		void from_setting(const Setting &set){
			Controller::from_setting(set);
			UI::get(yr,set,"yr",UI::compulsory);
			UI::get(b,set,"b",UI::compulsory);
			validate();
		}
		void validate() {	}
};
UIREGISTER(exact_ctrl);

/*! adaptive CE controller for \f$ y_{t+1} = y_{t} + b u_{t} + \sigma e_t \f$

From \f$ E\{y_{ref}\} = E\{ y_{t} + b u_{t} \} \f$, the optimal control action is:
\f[ u_t = (y_{ref} - y_{t})/b \f]

Using certainty equivalent approach, i.e. \f$ b= E(b|y_t, u_t, \ldots )\f$, which is 
estimated using ARX model. This class is not truly appropriate since it estimates also \f$ \sigma \f$


*/

class ce_ctrl: public exact_ctrl {
	protected:
		//! estimator
		ARX est;
		//! vector of "dyad" for estimation
		vec Psi;
		//! remember last action - for estimation
		double utm;
		//! remember last input - for estimation
		double ytm;
	public:
		//! expected input is yt
		virtual void adapt(const vec &yt){
			//prepare data for ARX
			Psi(0) = yt(0)-ytm;
			Psi(1) = utm;
			// do estimation
			est.bayes(Psi);
			// save estimate b = E(b);
			b = est.posterior().mean()(0); // first
		}
		//! same as exact control, but stores ut as utm
		vec ctrlaction(const vec &yt){
			vec ut=exact_ctrl::ctrlaction(yt);
			// remember last ut for estimation
			utm=ut(0);
			ytm=yt(0);
			return ut;
		}
		void from_setting(const Setting &set){
			Controller::from_setting(set);
			UI::get(yr,set,"yr",UI::compulsory);
			double b0;
			UI::get(b0,set,"b0",UI::compulsory);
			double P0;
			UI::get(P0,set,"P0",UI::compulsory);
			
			mat V0(2,2);
			V0(0,0) = 0.1;
			V0(1,0) = b0;
			V0(0,1) = b0;
			V0(1,1) = P0;
			
			est.set_statistics(1,V0);
			est.set_constant(false);
			validate();
		}
		void validate(){
			Psi = zeros(2);
			LIDs = zeros_i(1);
		}
		virtual void log_add ( logger &L, const string &name = "" ) {
			LIDs ( 0 ) = L.add ( RV("bhat",1) );
		}
		virtual void logit ( logger &L ) {
			L.logit ( LIDs ( 0 ), b );
		}
		
};
UIREGISTER(ce_ctrl);

#endif //ARX_CTRL