#include <stat/merger.h>

#ifdef MEX
#include <itpp/itmex.h>
#include <mex/mex_logger.h>
#include <mex/mex_parser.h>
#include <mex/config2mxstruct.h>
#endif

/*! \file
\brief Merging static pdfs
Merger is a program/mex-function for static merging of given pdfs on given support.

In command line, it expected config file with the following structure:
\code
Sources = (... any epdf or mpdf  ...);
Support = {grid=([low1d,high1d],[low2d,high2d],...);
		   nbins=[bins1d,bins2d,... ]};
// OR
Support = {pdf={class='epdf',...};
		   nsamples=2};
Merger = {class="merger_base",...}
\endcode

In mex-file, the above structures (Sources,Support,Merger) are input arguments. Type >>merger for help.
*/

using namespace bdm;

#ifdef MEX
void mexFunction (int n_output, mxArray *output[], int n_input, const mxArray *input[])
{
	// Check the number of inputs and output arguments
	if (n_input != 3) mexErrMsgTxt ("Usage:\n"
		                            "result=merger(sources, support, merger)\n"
		                            "  sources= { struct('class','epdf'),... };  % cell of pdfs (epdfs or mpdfs) to be merged,\n"
		                            "  support= struct(\n"
		                            "           grid    = {[dim1_start,dim1_end], [dim2_start, dim2_end]...}  %support boundary \n"
		                            "           nbins   = [bins_in_dim1, bins_in_dim2,...]                    %fixed \n"
		                            "         === OR ==\n"
		                            "           pdf     = struct('class','epdf'); % pdf to draw samples from\n"
		                            "           nsamples= 100;                    % number of samples\n"
		                            "           );\n"
		                            "        If all elements are present,  (grid,nbins) is used;\n"
		                            "  merger = struct('class','merger_*');       % object to be used for merging,\n\n"
		                            "see documentation of classes epdf, mpdf, merger_base and their offsprings in BDM.");

	// LOAD CONFIG
	UImxArray Cfg;
	Cfg.addList (input[0], "Sources");
	Cfg.addGroup (input[1], "Support");
	Cfg.addGroup (input[2], "Merger");

	//DBG
	Cfg.writeFile ("merger.cfg");
#else
int main() {
	UIFile Cfg ("merger.cfg");
#endif
	// Sources
	Array<mpdf*> Sources;
	//abuse Mer to store sources
	Setting& _Sources = Cfg.lookup ("Sources");
	int Slen = _Sources.getLength();
	Sources.set_size (Slen);
	for (int i = 0; i < Slen; i++) {
		try {
			mpdf* mtmp = UI::build<mpdf> (_Sources, i);
			Sources (i) = mtmp;
		} catch (UIException) {
			// it is not mpdf - see if it is epdf
			try {
				shared_ptr<epdf> etmp = UI::build<epdf> (_Sources, i);
				if (etmp) {
					Sources (i) = new mepdf (etmp);
				}
			} catch (UIException e) {
				it_error ("No mpdfs or epdfs found! " + string (e.what()));
			} catch (std::exception e) {
				it_error ("Error in UI at " + _Sources[i].getPath());
			}
		} catch (std::exception e) {
			it_error ("Error in UI at " + _Sources[i].getPath());
		}
	}

	merger_base* Merger = UI::build<merger_base> (Cfg, "Merger");

	// Support
	Setting & _Supp = Cfg.lookup ("Support");

	if (_Supp.exists ("grid") &&  _Supp.exists ("nbins")) {
		Array<vec> bounds (0);
		UI::get (bounds, _Supp, "grid");
		ivec nbins (0);
		UI::get (nbins, _Supp, "nbins");
		Merger->set_support (bounds, nbins);

	} else {
		if (_Supp.exists ("pdf") &&  _Supp.exists ("nsamples")) {
			epdf *g0 = UI::build<epdf> (_Supp, "pdf");
			int npoints = 100;
			_Supp.lookupValue ("nsamples", npoints);
			Merger->set_support (*g0, npoints);
			delete g0;
		} else it_error ("Use either [grid,nbins] or [pdf,nsamples].");
	}
// COMPUTE RESULTS
	Merger->set_sources (Sources, true); // takes care of deletion of sources
	Merger->merge();

// save results
	Array<vec> source_vals (Sources.length());
	for (int i = 0;i < Sources.length();i++) {
		datalink_m2e dl;
		dl.set_connection (Sources (i)->_rv(), Sources (i)->_rvc(), Merger->_rv());

		vec ll (Merger->_Smp()._samples().length());
		for (int j = 0; j < Merger->_Smp()._samples().length(); j++) {
			vec dt = dl.pushdown (Merger->_Smp()._samples() (j));
			vec dtc = dl.get_cond (Merger->_Smp()._samples() (j));
			ll (j) = Sources (i)->evallogcond (dt, dtc);
		}

		vec sll = exp (ll);

		source_vals (i) = sll / sum (sll);
	}

	merger_mix* MerMix=dynamic_cast<merger_mix*>(Merger);
	vec mix_val;
			
	if (MerMix){
		vec ll (Merger->_Smp()._samples().length());
		for (int j = 0; j < Merger->_Smp()._samples().length(); j++) {
			ll (j) = Merger->evallog (Merger->_Smp()._samples() (j));
		}
	
		vec sll = exp (ll);
	
		mix_val = sll / sum (sll);
	}

#ifdef MEX
	mxArray* tmp ;
	// Save results
	if (n_output > 0) {
		tmp = mxCreateStructMatrix (1, 1, 0, NULL);
		//support
		Array<vec> &samples = Merger->_Smp()._samples();
		if (samples.size() > 0) {
			mxArray* fld = mxCreateDoubleMatrix (samples (0).length(), samples.size(), mxREAL);
			Arrayvec2mxArray (samples, fld);
			mxReplaceFieldNM (tmp, "support", fld);
		}

		//weights
		vec &w = Merger->_Smp()._w();
		mxArray* fldw = mxCreateDoubleMatrix (1, w.length(), mxREAL);
		vec2mxArray (w, fldw);
		mxReplaceFieldNM (tmp, "weights", fldw);

		//mixture values
		if (mix_val.length()>0){
			mxArray* fldm = mxCreateDoubleMatrix (1, w.length(), mxREAL);
			vec2mxArray (mix_val, fldm);
			mxReplaceFieldNM (tmp, "mix", fldm);
		}
		// sources
		char srcstr[20];
		for (int i = 0;i < Sources.length();i++) {
			sprintf (srcstr, "source%d", i + 1);
			mxArray* fldw = mxCreateDoubleMatrix (1, source_vals (i).length(), mxREAL);
			vec2mxArray (source_vals (i), fldw);
			mxReplaceFieldNM (tmp, srcstr, fldw);
		}

		output[0] = tmp;
	}
#endif
}