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

namespace bdm {

MergerDiscrete::MergerDiscrete ( const Array<shared_ptr<pdf> > &S ) :
    Npoints ( 0 ), DBG ( false ), dbg_file ( 0 ) {
    set_sources ( S );
}


void MergerDiscrete::set_sources ( const Array<shared_ptr<pdf> > &PartSources ) {
    part_sources = PartSources;
    Nsources = part_sources.length();
    //set sizes
    sources.set_size ( Nsources);
	dls.set_size ( Nsources );
	rvzs.set_size ( Nsources );
	zdls.set_size ( Nsources );

    RV rv = get_composite_rv ( part_sources, /* checkoverlap = */ false );

    RV rvc;
    // Extend rv by rvc!
    for ( int i = 0; i < sources.length(); i++ ) {
        RV rvx = part_sources ( i )->_rvc().subt ( rv );
        rvc.add ( rvx ); // add rv to common rvc
    }

    // join rv and rvc - see descriprion
    rv.add ( rvc );
	
    // create links between sources and common rv
    RV xytmp;
    for ( int i = 0; i < part_sources.length(); i++ ) {		
        //Establich connection between pdfs and merger
        dls ( i ) = new datalink_m2e;
		dls ( i )->set_connection ( part_sources ( i )->_rv(), part_sources( i )->_rvc(), rv );

        // find out what is missing in each pdf
        xytmp = part_sources ( i )->_rv();
		xytmp.add ( part_sources ( i )->_rvc() );
        // z_i = common_rv-xy
        rvzs ( i ) = rv.subt ( xytmp );
        //establish connection between extension (z_i|x,y)s and common rv
        zdls ( i ) = new datalink_m2e;
        zdls ( i )->set_connection ( rvzs ( i ), xytmp, rv ) ;
    };
	
	// 
	merger().set_rv(rv);
}

void MergerDiscrete::set_support ( rectangular_support &Sup ) {
    Npoints = Sup.points();
    eSmp.set_parameters ( Npoints, false );
    Array<vec> &samples = eSmp._samples();
    eSmp._w() = ones ( Npoints ) / Npoints; //unifrom size of bins
    //set samples
    samples ( 0 ) = Sup.first_vec();
    for ( int j = 1; j < Npoints; j++ ) {
        samples ( j ) = Sup.next_vec();
    }
    eSmp.validate();
}

void MergerDiscrete::merge () {
    validate();

    //check if sources overlap:
    bool OK = true;
    for ( int i = 0; i < part_sources.length(); i++ ) {
        OK &= ( rvzs ( i )._dsize() == 0 ); // z_i is empty
        OK &= ( part_sources ( i )->_rvc()._dsize() == 0 ); // y_i is empty
    }

    if ( OK ) {
        mat lW = zeros ( part_sources.length(), eSmp._w().length() );

        vec emptyvec ( 0 );
        for ( int i = 0; i < part_sources.length(); i++ ) {
            for ( int j = 0; j < eSmp._w().length(); j++ ) {
                lW ( i, j ) = part_sources ( i )->evallogcond ( eSmp._samples() ( j ), emptyvec );
            }
        }

        vec w_nn = merge_points ( lW );
        vec wtmp = exp ( w_nn - max ( w_nn ) );
        //renormalize
        eSmp._w() = wtmp / sum ( wtmp );
    } else {
        bdm_error ( "Sources are not compatible - use merger_mix" );
    }
}

vec MergerDiscrete::merge_points ( mat &lW ) {
    int nu = lW.rows();
    vec result;
    ivec indW;
    bool infexist = false;
    switch ( METHOD ) {
    case ARITHMETIC:
        result = log ( sum ( exp ( lW ) ) ); //ugly!
        break;
    case GEOMETRIC:
        result = sum ( lW ) / nu;
        break;
    case LOGNORMAL:
        vec sumlW = sum ( lW ) ;
        indW = find ( ( sumlW < inf ) & ( sumlW > -inf ) );
        infexist = ( indW.size() < lW.cols() );
        vec mu;
        vec lam;
        if ( infexist ) {
            mu = sumlW ( indW ) / nu; //mean of logs
            //
            mat validlW = lW.get_cols ( indW );
            lam = sum ( pow ( validlW - outer_product ( ones ( validlW.rows() ), mu ), 2 ) );
        } else {
            mu = sum ( lW ) / nu; //mean of logs
            lam = sum ( pow ( lW - outer_product ( ones ( lW.rows() ), mu ), 2 ) );
        }
        //
        double coef = 0.0;
        vec sq2bl = sqrt ( 2 * beta * lam ); //this term is everywhere
        switch ( nu ) {
        case 2:
            coef = ( 1 - 0.5 * sqrt ( ( 4.0 * beta - 3.0 ) / beta ) );
            result = coef * sq2bl + mu ;
            break;
            // case 4: == can be done similar to case 2 - is it worth it???
        default: // see accompanying document merge_lognorm_derivation.lyx
            coef = sqrt ( 1 - ( nu + 1 ) / ( 2 * beta * nu ) );
            result = log ( besselk ( ( nu - 3 ) / 2, sq2bl * coef ) ) - log ( besselk ( ( nu - 3 ) / 2, sq2bl ) ) + mu;
            break;
        }
        break;
    }
    if ( infexist ) {
        vec tmp = -inf * ones ( lW.cols() );
        set_subvector ( tmp, indW, result );
        return tmp;
    } else {
        return result;
    }
}

void MergerDiscrete::from_setting ( const Setting& set ) {
    // get support
    // find which method to use
    string meth_str;
    UI::get( meth_str, set, "method", UI::compulsory );
    if ( meth_str == "arithmetic" )
        set_method ( ARITHMETIC );
    else if ( meth_str == "geometric" )
        set_method ( GEOMETRIC );
    else if ( meth_str ==  "lognormal" ) {
        set_method ( LOGNORMAL );
        UI::get(beta, set, "beta", UI::compulsory );
    }


    string dbg_filename;
    if ( UI::get ( dbg_filename, set, "dbg_file" ) )
        set_debug_file( dbg_filename );

}

void MergerDiscrete::to_setting  (Setting  &set) const {
    UI::save( METHOD, set, "method");

    if( METHOD == LOGNORMAL )
        UI::save (beta, set, "beta" );

    if( DBG )
        UI::save ( dbg_file->get_fname(), set, "dbg_file" );
}

void MergerDiscrete::validate() {
//        bdm_assert ( eSmp._w().length() > 0, "Empty support, use set_support()." );
//        bdm_assert ( dim == eSmp._samples() ( 0 ).length(), "Support points and rv are not compatible!" );
    MergerBase::validate();
    //bdm_assert ( merger().isnamed(), "mergers must be named" );
}

void merger_mix::merge ( ) {
	int dim = eSmp.dimension();
    if(Npoints<1) {
        set_support(enorm<fsqmat>(zeros(dim), eye(dim)), 1000);
    }

    bdm_assert(Npoints>0,"No points in support");
    bdm_assert(Nsources>0,"No Sources");

    Array<vec> &Smp = eSmp._samples(); //aux
    vec &w = eSmp._w(); //aux

    mat Smp_ex = ones ( dim, Npoints ); // Extended samples for the ARX model
    for ( int i = 0; i < Npoints; i++ ) {
        set_col_part ( Smp_ex, i, Smp ( i ) );
    }

    if ( DBG )    *dbg_file << Name ( "Smp_0" ) << Smp_ex;

    // Stuff for merging
    vec lw_src ( Npoints );        // weights of the ith source
    vec lw_mix ( Npoints );        // weights of the approximating mixture
    vec lw ( Npoints );            // tmp
    mat lW = zeros ( Nsources, Npoints ); // array of weights of all part_sources
    vec vec0 ( 0 );

    //initialize importance weights
    lw_mix = 1.0; // assuming uniform grid density -- otherwise

    // Initial component in the mixture model
    mat V0 = 1e-8 * eye ( dim + 1 );
    ARX A0;
    A0.set_statistics ( dim, V0 ); //initial guess of Mix:
    A0.validate();

    Mix.init ( &A0, Smp_ex, Ncoms );
    //Preserve initial mixture for repetitive estimation via flattening
    MixEF Mix_init ( Mix );

    // ============= MAIN LOOP ==================
    bool converged = false;
    int niter = 0;
    char dbg_str[100];

	emix *Mpred = Mix.epredictor ( );
    vec Mix_pdf ( Npoints );
    while ( !converged ) {
        //Re-estimate Mix
        //Re-Initialize Mixture model
        Mix.flatten ( &Mix_init , 1.0);
        Mix.bayes_batch_weighted ( Smp_ex, empty_vec, w*Npoints );
        delete Mpred;
        Mpred = Mix.epredictor ( ); // Allocation => must be deleted at the end!!
        Mpred->set_rv ( merger()._rv() ); //the predictor predicts rv of this merger

        // This will be active only later in iterations!!!
        if ( 1. / sum_sqr ( w ) < effss_coef*Npoints ) {
            // Generate new samples
            eSmp.set_samples ( Mpred );
            for ( int i = 0; i < Npoints; i++ ) {
                //////////// !!!!!!!!!!!!!
                //if ( Smp ( i ) ( 2 ) <0 ) {Smp ( i ) ( 2 ) = 0.01; }
                set_col_part ( Smp_ex, i, Smp ( i ) );
                //Importance of the mixture
                //lw_mix ( i ) =Mix.logpred (Smp_ex.get_col(i) );
                lw_mix ( i ) = Mpred->evallog ( Smp ( i ) );
            }
            if ( DBG ) {
                cout << "Resampling =" << 1. / sum_sqr ( w ) << endl;
                cout << Mix.posterior().mean() << endl;
                cout << sum ( Smp_ex, 2 ) / Npoints << endl;
                cout << Smp_ex*Smp_ex.T() / Npoints << endl;
            }
        }
        if ( DBG ) {
            sprintf ( dbg_str, "Mpred_mean%d", niter );
            *dbg_file << Name ( dbg_str ) << Mpred->mean();
            sprintf ( dbg_str, "Mpred_var%d", niter );
            *dbg_file << Name ( dbg_str ) << Mpred->variance();
            sprintf ( dbg_str, "Mpred_cov%d", niter );
            *dbg_file << Name ( dbg_str ) << Mpred->covariance();


            sprintf ( dbg_str, "pdf%d", niter );
            for ( int i = 0; i < Npoints; i++ ) {
                Mix_pdf ( i ) = Mix.logpred ( Smp_ex.get_col ( i ), empty_vec );
            }
            *dbg_file << Name ( dbg_str ) << Mix_pdf;

            sprintf ( dbg_str, "Smp%d", niter );
            *dbg_file << Name ( dbg_str ) << Smp_ex;

        }
        //Importace weighting
        for ( int i = 0; i < part_sources.length(); i++ ) {
            lw_src = 0.0;
            //======== Same RVs ===========
            //Split according to dependency in rvs
            if ( part_sources ( i )->dimension() == dim ) {
                // no need for conditioning or marginalization
                lw_src = part_sources ( i )->evallogcond_mat ( Smp , vec ( 0 ) );
            } else {
                // compute likelihood of marginal on the conditional variable
                if ( part_sources ( i )->dimensionc() > 0 ) {
                    // Make marginal on rvc_i
                    shared_ptr<epdf> tmp_marg = Mpred->marginal ( part_sources ( i )->_rvc() );
                    //compute vector of lw_src
                    for ( int k = 0; k < Npoints; k++ ) {
                        // Here val of tmp_marg = cond of pdfs(i) ==> calling dls->get_cond
                        lw_src ( k ) += tmp_marg->evallog ( dls ( i )->get_cond ( Smp ( k ) ) );
                    }

//                     sprintf ( str,"marg%d",niter );
//                     *dbg << Name ( str ) << lw_src;

                }
                // Compute likelihood of the missing variable
                if ( dim > ( part_sources ( i )->dimension() + part_sources ( i )->dimensionc() ) ) {
                    ///////////////
                    // There are variales unknown to pdfs(i) : rvzs
                    shared_ptr<pdf> tmp_cond = Mpred->condition ( rvzs ( i ) );
                    // Compute likelihood
                    vec lw_dbg = lw_src;
                    for ( int k = 0; k < Npoints; k++ ) {
                        lw_src ( k ) += log (
                                            tmp_cond->evallogcond (
                                                zdls ( i )->pushdown ( Smp ( k ) ),
                                                zdls ( i )->get_cond ( Smp ( k ) ) ) );
                        if ( !std::isfinite ( lw_src ( k ) ) ) {
                            lw_src ( k ) = -1e16;
                            cout << "!";
                        }
                    }
                }
                // Compute likelihood of the partial source
                for ( int k = 0; k < Npoints; k++ ) {
                    lw_src ( k ) += part_sources ( i )->evallogcond ( dls ( i )->pushdown ( Smp ( k ) ),
                                    dls ( i )->get_cond ( Smp ( k ) ) );
                }

            }

            lW.set_row ( i, lw_src ); // do not divide by mix
        }
        lw = MergerDiscrete::merge_points ( lW ); //merge

        //Importance weighting
        lw -=  lw_mix; // hoping that it is not numerically sensitive...
        w = exp ( lw - max ( lw ) );

        //renormalize
        double sumw = sum ( w );
        if ( std::isfinite ( sumw ) ) {
            w = w / sumw;
        } else {
            it_file itf ( "merg_err.it" );
            itf << Name ( "w" ) << w;
        }

        if ( DBG ) {
            sprintf ( dbg_str, "lW%d", niter );
            *dbg_file << Name ( dbg_str ) << lW;
            sprintf ( dbg_str, "w%d", niter );
            *dbg_file << Name ( dbg_str ) << w;
            sprintf ( dbg_str, "lw_m%d", niter );
            *dbg_file << Name ( dbg_str ) << lw_mix;
        }
        // ==== stopping rule ===
        niter++;
        converged = ( niter > stop_niter );
    }
    delete Mpred;
}

void merger_mix::from_setting ( const Setting& set ) {
    MergerDiscrete::from_setting ( set );
    Ncoms=DFLT_Ncoms;
    UI::get( Ncoms, set, "ncoms", UI::optional );
    effss_coef=DFLT_effss_coef;
    UI::get (effss_coef , set,  "effss_coef", UI::optional);
    stop_niter=10;
    UI::get ( stop_niter, set,"stop_niter", UI::optional );
}

void merger_mix::to_setting  (Setting  &set) const  {
    MergerDiscrete::to_setting(set);
    UI::save( Ncoms, set, "ncoms");
    UI::save (effss_coef , set,  "effss_coef");
    UI::save ( stop_niter, set,"stop_niter");
}

void merger_mix::validate() {
    MergerDiscrete::validate();
    bdm_assert(Ncoms>0,"Ncoms too small");
}

// DEFAULTS FOR MERGER_MIX
const int merger_mix::DFLT_Ncoms = 10;
const double merger_mix::DFLT_effss_coef = 0.9;

}
