#include "epdf_harness.h"
#include "base/bdmbase.h"
#include "base/user_info.h"
#include "stat/exp_family.h"
#include "mat_checks.h"
#include "test_util.h"
#include "UnitTest++.h"
#include <memory>

namespace bdm {

void epdf_harness::test_config ( const char *config_file_name ) {
    RV::clear_all();

    UIFile in ( config_file_name );
    Array<shared_ptr<epdf_harness> > input;
    UI::get ( input, in, "data", UI::compulsory );
    int sz = input.size();
    CHECK ( sz > 0 );
    for ( int i = 0; i < sz; ++i ) {
        input ( i )->test ( config_file_name, i );
    }
}

void epdf_harness::from_setting ( const Setting &set ) {
    hepdf = UI::build<epdf> ( set, "epdf", UI::compulsory );
    UI::get ( mean, set, "mean", UI::compulsory );
    UI::get ( variance, set, "variance", UI::compulsory );

    support = UI::build<support_base> ( set, "support", UI::optional );
    UI::get ( nsamples, set, "nsamples", UI::optional );
    UI::get ( R, set, "R", UI::optional );

    mrv = UI::build<RV> ( set, "marginal_rv", UI::optional );

    UI::get ( tolerance, set, "tolerance", UI::optional );
}

void epdf_harness::test ( const char *config_name, int idx ) {
    CurrentContext cc ( config_name, idx );

    CHECK_CLOSE_EX ( mean, hepdf->mean(), tolerance );
    CHECK_CLOSE_EX ( variance, hepdf->variance(), tolerance );

    if ( support ) { // support is given
        grid_fnc ep_disc;
        ep_disc.set_support ( support );
        ep_disc.set_values ( *hepdf );
        // ep_disc is discretized at support points

        CHECK_CLOSE ( 1.0,  ep_disc.integrate(), 0.01 );

        vec pdf = ep_disc._values();
        pdf /= sum ( pdf ); // normalize

        vec mea = pdf ( 0 ) * support->first_vec();
        mat Remp = pdf ( 0 ) * outer_product ( support->act_vec(), support->act_vec() );

        // run through all points
        for ( int i = 1; i < support->points(); i++ ) {
            mea += pdf ( i ) * support->next_vec();
            Remp += pdf ( i ) * outer_product ( support->act_vec(), support->act_vec() );
        }
        CHECK_CLOSE ( mean, mea, tolerance );
        CHECK_CLOSE ( R, Remp - outer_product ( mea, mea ), tolerance );
    }

    if ( variance.length() > 0 ) {
        check_sample_mean();
    }
    if ( R.rows() > 0 ) {
        check_covariance();
    }

    if ( mrv ) {
        RV crv = hepdf->_rv().subt ( *mrv );
        epdf_ptr m = hepdf->marginal ( *mrv );
        pdf_ptr c = hepdf->condition ( crv );

        pdf_array aa ( 2 );
        aa ( 0 ) = c;
        aa ( 1 ) = m;
        mprod mEp ( aa );

        check_cond_mean ( mEp );
        if ( R.rows() > 0 ) {
            check_cond_covariance ( mEp );
        }

        // test of pdflog at zero
        vec zero ( 0 );
        vec zeron = zeros ( hepdf->dimension() );

        double lpz = hepdf->evallog ( zeron );
        double lpzc = mEp.evallogcond ( zeron, zero );
        CHECK_CLOSE_EX ( lpz, lpzc, tolerance );

        vec lpzv ( 1 );
        lpzv ( 0 ) = lpz;

        mat zero1n ( hepdf->dimension(), 1 );
        for ( int i = 0; i < zero1n.rows(); ++i ) {
            zero1n ( i, 0 ) = 0;
        }

        vec lpzv_act = hepdf->evallog_mat ( zero1n );
        CHECK_CLOSE_EX ( lpzv, lpzv_act, tolerance );

        Array<vec> zeroa ( 3 );
        lpzv = vec ( zeroa.size() );
        for ( int i = 0; i < zeroa.size(); ++i ) {
            zeroa ( i ) = zeron;
            lpzv ( i ) = lpz;
        }

        lpzv_act = hepdf->evallog_mat ( zeroa );
        CHECK_CLOSE_EX ( lpzv, lpzv_act, tolerance );
    }
}


void epdf_harness::check_sample_mean() {
    vec delta = make_close_tolerance ( variance, nsamples );

    int tc = 0;
    Array<vec> actual ( CurrentContext::max_trial_count );
    do {
        mat smp = hepdf->sample_mat ( nsamples );
        vec emu = smp * ones ( nsamples ) / nsamples;
        actual ( tc ) = emu;
        ++tc;
    } while ( ( tc < CurrentContext::max_trial_count ) &&
              !UnitTest::AreClose ( mean, actual ( tc - 1 ), delta ) );

    if ( ( tc == CurrentContext::max_trial_count ) &&
            ( !UnitTest::AreClose ( mean, actual ( CurrentContext::max_trial_count - 1 ), delta ) ) ) {
        UnitTest::MemoryOutStream stream;
        stream << CurrentContext::format_context ( __LINE__ ) << "expected " << mean << " +/- " << delta << " but was " << actual;

        UnitTest::TestDetails details ( *UnitTest::CurrentTest::Details(), 0, false );

        UnitTest::CurrentTest::Results()->OnTestFailure ( details, stream.GetText() );
    }
}

void epdf_harness::check_covariance() {
    int tc = 0;
    Array<mat> actual ( CurrentContext::max_trial_count );
    do {
        mat smp = hepdf->sample_mat ( nsamples );
        vec emu = smp * ones ( nsamples ) / nsamples;
        mat er = ( smp * smp.T() ) / nsamples - outer_product ( emu, emu );
        actual ( tc ) = er;
        ++tc;
    } while ( ( tc < CurrentContext::max_trial_count ) &&
              !UnitTest::AreClose ( R, actual ( tc - 1 ), tolerance ) );

    if ( ( tc == CurrentContext::max_trial_count ) &&
            ( !UnitTest::AreClose ( R, actual ( CurrentContext::max_trial_count - 1 ), tolerance ) ) ) {
        UnitTest::MemoryOutStream stream;
        stream << CurrentContext::format_context ( __LINE__ ) << "expected " << R << " +/- " << tolerance << " but was " << actual;

        UnitTest::TestDetails details ( *UnitTest::CurrentTest::Details(), 0, false );

        UnitTest::CurrentTest::Results()->OnTestFailure ( details, stream.GetText() );
    }
}

void epdf_harness::check_cond_mean ( mprod &mep ) {
    vec delta = make_close_tolerance ( variance, nsamples );

    int tc = 0;
    Array<vec> actual ( CurrentContext::max_trial_count );
    do {
        mat smp = mep.samplecond_mat ( vec ( 0 ), nsamples );
        vec emu = sum ( smp, 2 ) / nsamples;
        actual ( tc ) = emu;
        ++tc;
    } while ( ( tc < CurrentContext::max_trial_count ) &&
              !UnitTest::AreClose ( mean, actual ( tc - 1 ), delta ) );
    if ( ( tc == CurrentContext::max_trial_count ) &&
            ( !UnitTest::AreClose ( mean, actual ( CurrentContext::max_trial_count - 1 ), delta ) ) ) {
        UnitTest::MemoryOutStream stream;
        stream << CurrentContext::format_context ( __LINE__ ) << "expected " << mean << " +/- " << delta << " but was " << actual;

        UnitTest::TestDetails details ( *UnitTest::CurrentTest::Details(), 0, false );

        UnitTest::CurrentTest::Results()->OnTestFailure ( details, stream.GetText() );
    }
}

void epdf_harness::check_cond_covariance ( mprod &mep ) {
    int tc = 0;
    Array<mat> actual ( CurrentContext::max_trial_count );
    do {
        mat smp = mep.samplecond_mat ( vec ( 0 ), nsamples );
        vec emu = sum ( smp, 2 ) / nsamples;
        mat er = ( smp * smp.T() ) / nsamples - outer_product ( emu, emu );
        actual ( tc ) = er;
        ++tc;
    } while ( ( tc < CurrentContext::max_trial_count ) &&
              !UnitTest::AreClose ( R, actual ( tc - 1 ), tolerance ) );
    if ( ( tc == CurrentContext::max_trial_count ) &&
            ( !UnitTest::AreClose ( R, actual ( CurrentContext::max_trial_count - 1 ), tolerance ) ) ) {
        UnitTest::MemoryOutStream stream;
        stream << CurrentContext::format_context ( __LINE__ ) << "expected " << mean << " +/- " << tolerance << " but was " << actual;

        UnitTest::TestDetails details ( *UnitTest::CurrentTest::Details(), 0, false );

        UnitTest::CurrentTest::Results()->OnTestFailure ( details, stream.GetText() );
    }
}

}
