#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)const{
			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 &dt){
			//prepare data for ARX
			// do estimation
			est.bayes(vec_1(dt(0)-ytm),vec_1(dt(1)));
			ytm = dt(0);
			// save estimate b = E(b);
			b = est.posterior().mean()(0); // first
		}
		//! same as exact control, but stores ut as utm
		vec ctrlaction(const vec &dt)const{
			vec ut(1);
			mat P = est.posterior().variance();
			ut(0) =(yr - dt(0))/(b + (b<0 ? P(0,0) : -P(0,0)));;
			// remember last ut for estimation
			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/6;
			
			est.set_statistics(1,V0,6);
			est.set_constant(false);
			est.validate();
			validate();
		}
		void validate(){
			Psi = zeros(2);
		}
		void log_register( logger& L, const string& prefix ){
			est.set_log_level(1);
			est.log_register(L,"");
		};
		void log_write() const{
			est.log_write();
		}
};
UIREGISTER(ce_ctrl);

#endif //ARX_CTRL