#include "estim/arx.h"
#include "../mat_checks.h"

using namespace bdm;

TEST ( arx_elem_stress ) {
    // Setup model : ARX for 1D Gaussian
    //Test constructor
    mat V0 = 0.00001 * eye ( 2 );
    V0 ( 0, 0 ) = 0.1; //
    ARX Ar;
    Ar.set_statistics ( 1, V0, -1.0 );
    Ar.set_constant ( true );
    Ar.validate();

    mat mu ( 1, 1 );
    mat R ( 1, 1 );
    Ar.posterior().mean_mat ( mu, R );
    cout << "Prior moments: mu=" << mu << ", R=" << R << endl;

    int ndat = 200;
    vec smp = randn ( ndat );
    //
    mat Smp = ones ( 2, ndat );
    Smp.set_row ( 0, smp );
    //
    Ar.bayes_batch ( Smp );
    // Ar is now filled with estimates of N(0,1);
    cout << "Empirical moments: mu=" << sum ( smp ) / ndat << ", R=" << sum_sqr ( smp ) / ndat - pow ( sum ( smp ) / ndat, 2 ) << endl;
    Ar.posterior().mean_mat ( mu, R );
    cout << "Posterior moments: mu=" << mu << ", R=" << R << endl;

    //////// TEST prediction
    vec x = linspace ( -3.0, 3.0, 100 );
    double xstep = 6.0 / 100.0;
    mat X ( 1, 100 );
    mat X2 ( 2, 100 );
    X.set_row ( 0, x );
    X2.set_row ( 0, x );

    mlstudent* Ap = Ar.predictor_student();
    vec Ap_x = Ap->evallogcond_mat ( X, empty_vec );
    vec ll_x = Ar.logpred_mat ( X2 , empty_vec);

    cout << "normalize : " << xstep*sum ( exp ( Ap_x ) ) << endl;
    cout << "normalize : " << xstep*sum ( exp ( ll_x ) ) << endl;

    it_file it ( "arx_elem_test.it" );
    it << Name ( "Ap_x" ) << Ap_x;
    it << Name ( "ll_x" ) << ll_x;
}
