/*
Copyright 2013 Chipman, H.A., George, E.,I., McCulloch, R.E., and Pratola, M.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/
#include <iostream>
#include <fstream>
#include <vector>
#include <ctime>
#include <sstream>

#include "rng.h"
#include "tree.h"
#include "info.h"
#include "funs.h"
#include "bd.h"
#include "mpi.h"

using std::cout;
using std::endl;

int main(int argc, char** argv)
{
	int myrank, processes;
	MPI_Init(&argc,&argv);
	MPI_Comm_rank(MPI_COMM_WORLD,&myrank);
	MPI_Comm_size(MPI_COMM_WORLD,&processes);

#ifndef SILENT
	cout << "\nMPI: node " << myrank << " of " << processes << " processes.";
	cout << "\n*****Into bart main\n";
#endif

	//Load the fitted trees and cutpoints
	size_t m,p,ndraws;
	xinfo xi;
	std::ifstream treef("trees.txt");
	treef >> xi; //load the cutpoints
	treef >> m;  //number of trees
	treef >> p;  //dimension of x's
	treef >> ndraws; //number of draws from the posterior that were saved.

	std::vector<std::vector<tree> > t(ndraws,std::vector<tree>(m));
	for(size_t i=0;i<ndraws;i++)
		for(size_t j=0;j<m;j++)
			treef >> t[i][j];


	//Load the prediction points for each of the nodes
	if(argc<2) {
		cout << "error, expecting prediction root filename as first argument\n"; 
		return 1;
	}
	double xtemp; //used to read in xpred without knowing how many there are
	char ptype='p'; //type of prediction, 'p' means mean of posterior predictive, 'd' means a single draw from
					//the posterior predictive
	if(argc>2) ptype=(char) *argv[2];
	dinfo dip;  //data information for prediction
	dip.n=0;
	std::vector<double> xp;  //prediction observations, stored like x
	std::stringstream xfss;
	std::string xfs;
	xfss << argv[1] << myrank+1 << ".txt";
	xfs=xfss.str();
	std::ifstream xpf(xfs.c_str()); //file to read x to predict at from for this node

#ifndef SILENT
	cout << myrank << ": loading xp's from " << xfs << endl;
#endif
	while(xpf >> xtemp)
		xp.push_back(xtemp);
	size_t np = xp.size()/p;
	if(xp.size() != np*p) {
		cout << "error, wrong number of elements in prediction data set\n";
		return 1;
	}
	if(np)
		dip.n=np; dip.p=p; dip.x=&xp[0]; dip.y=0; //there are no y's!

#ifndef SILENT
	cout << "\nx for prediction read in:\n";
	cout << myrank << ": n: " << dip.n << endl;
	if(dip.n) {
		cout << myrank << ": first row: " << dip.x[0] << " ... " << dip.x[p-1] << endl;
		cout << myrank << ":  last row: " << dip.x[(dip.n-1)*p] << " ... " << dip.x[dip.n*p-1] << endl;
	}
#endif

	//Do this nodes predictions
	double* ppredmean=0; //posterior mean for prediction
	double* fpredtemp=0; //temporary fit vector to compute prediction
	double* postp=0; //will hold the posterior mean over all the draws saved from the mcmc
	double* postp2=0; //will hold the posterior mean of E[f^2] over all the draws saved from the mcmc
	ppredmean = new double[dip.n];
  	fpredtemp = new double[dip.n];
	postp = new double[dip.n];
	postp2 = new double[dip.n];
   	for(size_t i=0;i<dip.n;i++) {
		ppredmean[i]=0.0;
		postp[i]=0.0;
		postp2[i]=0.0;
	}

	if(ptype=='p')
	{
#ifndef SILENT
		cout << "Getting sample mean and stdev of posterior predictive" << endl;
#endif
		for(size_t i=0;i<ndraws;i++) {
			for(size_t j=0;j<m;j++) {
				fit(t[i][j],xi,dip,fpredtemp);
				for(size_t k=0;k<dip.n;k++) ppredmean[k] += fpredtemp[k];
			}
			for(size_t k=0;k<dip.n;k++) {
				if(i>0)
				{
					postp[k] *= (int)i;
					postp[k] += ppredmean[k];
					postp[k] /= (int)(i+1);
					postp2[k] *= (int)i;
					postp2[k] += ppredmean[k]*ppredmean[k];
					postp2[k] /= (int)(i+1);
				}
				else
				{
					postp[k] = ppredmean[k];
					postp2[k] = ppredmean[k]*ppredmean[k];
				}
				ppredmean[k]=0.0;
			}
		}
	}
	else if(ptype=='d') //ptype='d', a single draw from posterior predictive
	{
#ifndef SILENT
		cout << "Getting a single draw from the posterior predictive" << endl;
#endif
		size_t draw=0;
		double genunif;
		uint seed;
		if(myrank==0) seed=(unsigned int)time(NULL);
		MPI_Bcast(&seed,1,MPI_UNSIGNED,0,MPI_COMM_WORLD);
		RNG gen(seed);
		genunif=gen.uniform();
#ifndef SILENT
		cout << myrank << ": seed=" << seed << " unif=" << genunif << endl;
#endif

		for(size_t i=0;i<ndraws;i++)
			if(((double)(i)/(double)(ndraws))<genunif)
				draw=i;

#ifndef SILENT
		cout << "drawing the i=" << draw << "th tree for posterior prediction" << endl;
#endif

		for(size_t j=0;j<m;j++) {
			fit(t[draw][j],xi,dip,fpredtemp);
			for(size_t k=0;k<dip.n;k++) ppredmean[k] += fpredtemp[k];
		}
		for(size_t k=0;k<dip.n;k++) {
			postp[k] = ppredmean[k];
			postp2[k] = ppredmean[k]*ppredmean[k];
		}

	}
	else
	{
		cout << "error, prediction type " << ptype << " is invalid\n";
		return 1;
	}

	//Write out this nodes predictions
	std::stringstream yfps;
	std::string yfs;
	yfps << "yp" << myrank << ".txt";
	yfs=yfps.str();
#ifndef SILENT
	cout << myrank << ": saving yp's to " << yfs << endl;
#endif
	std::ofstream ypf(yfs.c_str()); //file to save the predictions to for this node
	for(size_t i=0;i<dip.n;i++)
		ypf << postp[i] << endl;

	//Write out this nodes std deviations
	std::stringstream yfsd;
	std::string ysd;
	yfsd << "ysd" << myrank << ".txt";
	ysd=yfsd.str();
#ifndef SILENT
	cout << myrank <<": saving ysd's to " << ysd << endl;
#endif
	std::ofstream ysdf(ysd.c_str()); //file to save the std deviations to for this node
	for(size_t i=0;i<dip.n;i++)
		ysdf << sqrt(postp2[i]-postp[i]*postp[i]) << endl;

	//cleanup
	delete[] ppredmean;
	delete[] fpredtemp;
	delete[] postp;
	delete[] postp2;

	MPI_Finalize();

	return 0;
} //end of main
