#include "arx.h"
namespace bdm {

void ARX::bayes_weighted ( const vec &yt, const vec &cond, const double w ) {
    bdm_assert_debug ( yt.length() == dimy, "BM::bayes yt is of size "+num2str(yt.length())+" expected dimension is "+num2str(dimy) );
    bdm_assert_debug ( cond.length() == rgrlen , "BM::bayes cond is of size "+num2str(cond.length())+" expected dimension is "+num2str(rgrlen) );

    BMEF::bayes_weighted(yt,cond,w); //potential discount scheduling
    double lnc;
    //cache
    ldmat &V = est._V();
    double &nu = est._nu();

    dyad.set_subvector ( 0, yt );
    if (cond.length()>0)
        dyad.set_subvector ( dimy, cond );
    // possible "1" is there from the beginning

    if ( frg < 1.0 ) {
        est.pow ( frg ); // multiply V and nu


        //stabilize
        ldmat V0 = alter_est._V(); //$ copy
        double &nu0 = alter_est._nu();

        V0 *= ( 1 - frg );
        V += V0; //stabilization
        nu += ( 1 - frg ) * nu0;

        // recompute loglikelihood of new "prior"
        if ( evalll ) {
            last_lognc = est.lognc();
        }
    }
    V.opupdt ( dyad, w );
    nu += w;

    // log(sqrt(2*pi)) = 0.91893853320467
    if ( evalll ) {
        lnc = est.lognc();
        ll = lnc - last_lognc - 0.91893853320467;
        last_lognc = lnc;
    }
}

double ARX::logpred ( const vec &yt, const vec &cond ) const {
    egiw pred ( est );
    ldmat &V = pred._V();
    double &nu = pred._nu();

    double lll;
    vec dyad_p = dyad;
    dyad_p.set_subvector ( 0, yt );
    dyad_p.set_subvector(dimy,cond);

    if ( frg < 1.0 ) {
        pred.pow ( frg );
        lll = pred.lognc();
    } else//should be save: last_lognc is changed only by bayes;
        if ( evalll ) {
            lll = last_lognc;
        } else {
            lll = pred.lognc();
        }

    V.opupdt ( dyad_p, 1.0 );
    nu += 1.0;
    // log(sqrt(2*pi)) = 0.91893853320467
    return pred.lognc() - lll - 0.91893853320467;
}

vec ARX::samplepred(const vec &cond) const{
	mat M;
	chmat R;
	est.sample_mat(M,R);
	
	vec tmp;
	if (dimc>0){
		if (have_constant){
			tmp = M*concat(cond,1.0);
		} else {
			tmp = M*cond;
		}
	} else {
		tmp = zeros(dimy);
	}
	tmp += R._Ch().T()*randn(dimy);
	return tmp;
}

void ARX::flatten ( const BMEF* B , double weight =1.0) {
    const ARX* A = dynamic_cast<const ARX*> ( B );
    // nu should be equal to B.nu
    est.pow ( A->posterior()._nu() / posterior()._nu() *weight);
    if ( evalll ) {
        last_lognc = est.lognc();
    }
}

ARX* ARX::_copy ( ) const {
    ARX* Tmp = new ARX ( *this );
    return Tmp;
}

void ARX::set_statistics ( const BMEF* B0 ) {
    const ARX* A0 = dynamic_cast<const ARX*> ( B0 );

    bdm_assert_debug ( dimension() == A0->dimension(), "Statistics of different dimensions" );
    set_statistics ( A0->dimensiony(), A0->posterior()._V(), A0->posterior()._nu() );
}

enorm<ldmat>* ARX::epredictor ( const vec &cond ) const {
    bdm_assert_debug ( cond.length() == rgrlen , "ARX::epredictor cond is of size "+num2str(cond.length())+" expected dimension is "+num2str(rgrlen) );

    mat mu ( dimy, posterior()._V().rows() - dimy );
    mat R ( dimy, dimy );

    vec ext_rgr;
    if (have_constant) {
        ext_rgr = concat(cond,vec_1(1.0));
    } else {
        ext_rgr = cond;
    }

    enorm<ldmat>* tmp;
    tmp = new enorm<ldmat> ( );
    //TODO: too hackish
    if ( yrv._dsize() > 0 ) {
    }

    est.mean_mat ( mu, R ); //mu =
    //correction for student-t  -- TODO check if correct!!
    R*=posterior()._nu()/(posterior()._nu()-2);
    if (mu.cols()>0) {// nonempty egiw
        mat p_mu = mu.T() * ext_rgr; 	//the result is one column
        tmp->set_parameters ( p_mu.get_col ( 0 ), ldmat ( R ) );
    } else {
        tmp->set_parameters ( zeros( R.rows() ), ldmat ( R ) );
    }
    if (dimy==yrv._dsize())
        tmp->set_rv(yrv);
    return tmp;
}

mlstudent* ARX::predictor_student ( ) const {
    const ldmat &V = posterior()._V();

    mat mu ( dimy, V.rows() - dimy );
    mat R ( dimy, dimy );
    mlstudent* tmp;
    tmp = new mlstudent ( );

    est.mean_mat ( mu, R ); //
    mu = mu.T();

    int end = V._L().rows() - 1;
    ldmat Lam ( V._L() ( dimy, end, dimy, end ), V._D() ( dimy, end ) );  //exp val of R


    if ( have_constant ) { // no constant term
        //Assume the constant term is the last one:
        if ( mu.cols() > 1 ) {
            tmp->set_parameters ( mu.get_cols ( 0, mu.cols() - 2 ), mu.get_col ( mu.cols() - 1 ), ldmat ( R ), Lam );
        } else {
            tmp->set_parameters ( mat ( dimy, dimc ), mu.get_col ( mu.cols() - 1 ), ldmat ( R ), Lam );
        }
    } else {
        // no constant term
        tmp->set_parameters ( mu, zeros ( dimy ), ldmat ( R ), Lam );
    }
    return tmp;
}



/*! \brief Return the best structure
@param Eg a copy of GiW density that is being examined
@param Eg0 a copy of prior GiW density before estimation
@param Egll likelihood of the current Eg
@param indices current indices
\return best likelihood in the structure below the given one
*/
double egiw_bestbelow ( egiw Eg, egiw Eg0, double Egll, ivec &indices ) { //parameter Eg is a copy!
    ldmat Vo = Eg._V(); //copy
    ldmat Vo0 = Eg._V(); //copy
    ldmat& Vp = Eg._V(); // pointer into Eg
    ldmat& Vp0 = Eg._V(); // pointer into Eg
    int end = Vp.rows() - 1;
    int i;
    mat Li;
    mat Li0;
    double maxll = Egll;
    double tmpll = Egll;
    double belll = Egll;

    ivec tmpindices;
    ivec maxindices = indices;


    cout << "bb:(" << indices << ") ll=" << Egll << endl;

    //try to remove only one rv
    for ( i = 0; i < end; i++ ) {
        //copy original
        Li = Vo._L();
        Li0 = Vo0._L();
        //remove stuff
        Li.del_col ( i + 1 );
        Li0.del_col ( i + 1 );
        Vp.ldform ( Li, Vo._D() );
        Vp0.ldform ( Li0, Vo0._D() );
        tmpll = Eg.lognc() - Eg0.lognc(); // likelihood is difference of norm. coefs.

        cout << "i=(" << i << ") ll=" << tmpll << endl;

        //
        if ( tmpll > Egll ) { //increase of the likelihood
            tmpindices = indices;
            tmpindices.del ( i );
            //search for a better match in this substructure
            belll = egiw_bestbelow ( Eg, Eg0, tmpll, tmpindices );
            if ( belll > maxll ) { //better match found
                maxll = belll;
                maxindices = tmpindices;
            }
        }
    }
    indices = maxindices;
    return maxll;
}

ivec ARX::structure_est ( const egiw &est0 ) {
    ivec ind = linspace ( 1, est.dimension() - 1 );
    egiw_bestbelow ( est, est0, est.lognc() - est0.lognc(), ind );
    return ind;
}



ivec ARX::structure_est_LT ( const egiw &est0 ) {
    //some stuff with beliefs etc.
    ivec belief = vec_1 ( 2 );        // default belief
    int nbest = 1;           // nbest: how many regressors are returned
    int nrep = 5;         // nrep: number of random repetions of structure estimation
    double lambda   = 0.9;
    int k = 2;

    Array<str_aux> o2;

    ivec ind = bdm::straux1(est._V(),est._nu(), est0._V(), est0._nu(), belief, nbest, nrep, lambda, k, o2);

    return ind;
}

void ARX::from_setting ( const Setting &set ) {
    BMEF::from_setting(set);

    UI::get (rgr, set, "rgr", UI::compulsory );

    dimy = yrv._dsize();
    bdm_assert(dimy>0,"ARX::yrv should not be empty");
    rgrlen = rgr._dsize();

    int constant;
    if ( !UI::get ( constant, set, "constant", UI::optional ) ) {
        have_constant = true;
    } else {
        have_constant = constant > 0;
    }
    dimc = rgrlen;
    rvc = rgr;

    //init
    shared_ptr<egiw> pri = UI::build<egiw> ( set, "prior", UI::optional );
    if (pri) {
        set_prior(pri.get());
    } else {
        shared_ptr<egiw> post = UI::build<egiw> ( set, "posterior", UI::optional );
        set_prior(post.get());
    }


    shared_ptr<egiw> alt = UI::build<egiw> ( set, "alternative", UI::optional );
    if ( alt ) {
        bdm_assert ( alt->_dimx() == dimy, "alternative is not compatible" );
        bdm_assert ( alt->_V().rows() == dimy + rgrlen + int(have_constant==true), "alternative is not compatible" );
        alter_est.set_parameters ( alt->_dimx(), alt->_V(), alt->_nu() );
        alter_est.validate();
    }
    // frg handled by BMEF

}

void ARX::set_prior(const epdf *pri) {
    const egiw * eg=dynamic_cast<const egiw*>(pri);
    if ( eg ) {
        bdm_assert ( eg->_dimx() == dimy, "prior is not compatible" );
        bdm_assert ( eg->_V().rows() == dimy + rgrlen + int(have_constant==true), "prior is not compatible" );
        est.set_parameters ( eg->_dimx(), eg->_V(), eg->_nu() );
        est.validate();
    } else {
        est.set_parameters ( dimy, zeros ( dimy + rgrlen +int(have_constant==true)) );
        set_prior_default ( est );
    }
    //check alternative
    if (alter_est.dimension()!=dimension()) {
        alter_est = est;
    }
}

void ARXpartialforg::bayes ( const vec &val, const vec &cond ) {
#define LOG2  0.69314718055995
    vec frg = cond.right(cond.length() - rgrlen);
    vec cond_rgr = cond.left(rgrlen);       // regression vector

    int dimV = est._V().cols();
    int nparams = dimV - 1;                 // number of parameters
    int nalternatives = 1 << nparams;    // number of alternatives

    // Permutation matrix
    mat perm_matrix = ones(nalternatives, nparams);
    int i, j, period, idx_from, idx_to, start, end;
    for(i = 0; i < nparams; i++) {
        idx_from = 1 << i;
        period   = ( idx_from << 1 );
        idx_to   = period - 1;
        j = 0;
        start = idx_from;
        end = idx_to;
        while ( start < nalternatives ) {
            perm_matrix.set_submatrix(start, end, i, i, 0);
            j++;
            start = idx_from + period * j;
            end = idx_to + period * j;
        }
    }

    // Array of egiws for approximation
    Array<egiw*> egiw_array(nalternatives + 1);
    // No. of conditioning rows in LD
    int nalternatives_cond, position;

    for(int i = 0; i < nalternatives; i++) {
        // vector defining alternatives
        vec vec_alt = perm_matrix.get_row(i);

        // full alternative or filtered
        if( sum(vec_alt) == vec_alt.length() )  {
            egiw_array(i) = &alter_est;
            continue;
        } else if( sum(vec_alt) == 0 ) {
            egiw_array(i) = &est;
            continue;
        }

        nalternatives_cond = (int) sum(vec_alt) + 1;
        ivec vec_perm(0);                   // permutation vector

        for(int j = 0; j < nparams; j++) {
            position = dimV - j - 2;
            if ( vec_alt(position) == 0 ) {
                vec_perm.ins(j, position + 1);
            }
            else {
                vec_perm.ins(0, position + 1);
            }
        }
        vec_perm.ins(0, 0);

        ldmat filt (est._V(), vec_perm);
        ldmat alt (alter_est._V(), vec_perm);

        mat tmpL(dimV, dimV);
        tmpL.set_rows( 0, alt._L().get_rows(0, nalternatives_cond - 1) );
        tmpL.set_rows( nalternatives_cond, filt._L().get_rows(nalternatives_cond, nparams) );

        vec tmpD(dimV);
        tmpD.set_subvector( 0, alt._D()(0, nalternatives_cond - 1) );
        tmpD.set_subvector( nalternatives_cond, filt._D()(nalternatives_cond, nparams) );

        ldmat tmpLD (tmpL, tmpD);

        vec_perm = sort_index(vec_perm);
        ldmat newLD (tmpLD, vec_perm);

        egiw_array(i) = new egiw(1, newLD, alter_est._nu());
    }

    // Approximation
    double sumVecCommon;	      	// frequently used term
    vec vecNu(nalternatives);		// vector of nus of components
    vec vecD(nalternatives);		// vector of LS reminders
    vec vecCommon(nalternatives);	// vector of common parts
    mat matVecsTheta;				// matrix whose rows are theta vects.

    for (i = 0; i < nalternatives; i++) {
        vecNu.shift_left( egiw_array(i)->_nu() );
        vecD.shift_left( egiw_array(i)->_V()._D()(0) );
        matVecsTheta.append_row( egiw_array(i)->est_theta()  );
    }

    vecCommon = elem_mult ( frg, elem_div(vecNu, vecD) );
    sumVecCommon = sum(vecCommon);

    // approximation of est. regr. coefficients
    vec aprEstTheta(nparams);
    aprEstTheta.zeros();
    for (i = 0; i < nalternatives; i++) {
        aprEstTheta +=  matVecsTheta.get_row(i) * vecCommon(i);
    }
    aprEstTheta /= sumVecCommon;

    // approximation of degr. of freedom
    double aprNu;
    double A = log( sumVecCommon );		// Term 'A' in equation
    for ( int i = 0; i < nalternatives; i++ ) {
        A += frg(i) * ( log( vecD(i) ) - psi( 0.5 * vecNu(i) ) );
    }
    aprNu = ( 1 + sqrt(1 + 4 * (A - LOG2)/3 ) ) / ( 2 * (A - LOG2) );

    // approximation of LS reminder D(0,0)
    double aprD = aprNu / sumVecCommon;

    // Aproximation of covariance of LS est.
    mat aprC = zeros(nparams, nparams);
    for ( int i = 0; i < nalternatives; i++ ) {
        aprC += egiw_array(i)->est_theta_cov().to_mat() * frg(i);
        vec tmp = ( matVecsTheta.get_row(i) - aprEstTheta );
        aprC += vecCommon(i) * outer_product( tmp, tmp);
    }

    // Construct GiW pdf
    ldmat aprCinv ( inv(aprC) );
    vec D = concat( aprD, aprCinv._D() );
    mat L = eye(dimV);
    L.set_submatrix(1, 0, aprCinv._L() * aprEstTheta);
    L.set_submatrix(1, 1, aprCinv._L());
    ldmat aprLD (L, D);
    est = egiw(1, aprLD, aprNu);

    if ( evalll ) {
        last_lognc = est.lognc();
    }

    // update
    ARX::bayes ( val, cond_rgr );
}

}