/*
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
   if(argc<3) {cout << "error need at least 2 arguments\n"; return 1;}
   std::stringstream xroot(argv[1]);  // this is the root of the x and y filenames for bart to read data
   std::stringstream yroot(argv[2]);  // in from.  The format of data files must be <root>#.txt. The number
                                  // represented by # is the node number for the parallel slave's, and 
                                  // as such # must be greater than 0.
   double xtemp,ytemp; //used to read in x and y without knowing how many there are.

   //--------------------------------------------------
   //random number generation
   uint seed=99;
	MPI_Bcast(&seed,1,MPI_UNSIGNED,0,MPI_COMM_WORLD);
   RNG gen(seed); //this one random number generator is used in all draws
#ifndef SILENT
	cout << "RNG seed set to " << seed << "\n";
#endif

   //--------------------------------------------------
   //read in data
   //read y NOTE: we assume y is already centered at 0.
	//IMPORTANT NOTES FOR MPI VERSION: 
	// allys has correct information on the root node only
	// miny and maxy have the min/max of the y's loaded on the each slave respectively, and the min/max of the whole
	// dataset on the root node
	// n is the size of the entire dataset, nn is the size of the dataset loaded on each slave respectively
	sinfo allys;
	double sy=0.0, sy2=0.0;
	unsigned int tempn=0;
	size_t n=0,nn=0;
	double miny=INFINITY;
	double maxy=-INFINITY;
	double myminy=INFINITY;
	double mymaxy=-INFINITY;
	std::vector<double> y;   //storage for y on the slaves
	if(myrank!=0)
	{
		std::stringstream yfss;
		std::string yfs;
//		yfss << "y" << myrank << ".txt";
		yfss << yroot.str() << myrank << ".txt";
		yfs=yfss.str();
#ifndef SILENT
		cout << myrank << ": loading y's from " << yfs << endl;
#endif
		std::ifstream yf(yfs.c_str()); //file to read y from for this slave
	
		while(yf >> ytemp) {
			y.push_back(ytemp);
			if(ytemp<myminy) myminy=ytemp;
			if(ytemp>mymaxy) mymaxy=ytemp;
			sy+=ytemp;
			sy2+=ytemp*ytemp;
		}
		tempn=y.size();
		nn=(size_t)tempn;
		if(nn<1) {
			cout << "error nn<1\n";
			MPI_Finalize();
			return 1;
		}
#ifndef SILENT
		cout << myrank << ": nn=" << nn << " sy=" << sy << " sy2=" << sy2 << " myminy=" << myminy << " mymaxy=" << mymaxy << endl;
#endif
		miny=myminy;
		maxy=mymaxy;
	}
	MPI_Allreduce(&sy,&allys.sy,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
	MPI_Allreduce(&sy2,&allys.sy2,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
	MPI_Allreduce(&tempn,&tempn,1,MPI_UNSIGNED,MPI_SUM,MPI_COMM_WORLD);
	MPI_Allreduce(&myminy,&miny,1,MPI_DOUBLE,MPI_MIN,MPI_COMM_WORLD);
	MPI_Allreduce(&mymaxy,&maxy,1,MPI_DOUBLE,MPI_MAX,MPI_COMM_WORLD);
	n=(size_t) tempn;
#ifndef SILENT
	cout << myrank << ": n=" << n << " nn=" << nn << " allys.sy=" << allys.sy << " allys.sy2=" << allys.sy2 << " miny=" << miny << " maxy=" << maxy << "\n";
#endif
	if(n<1) {
		cout << "error n<1\n";
		MPI_Finalize();
		return 1;
	}
	allys.n=n;
   double ybar = allys.sy/n; //sample mean
   double shat = sqrt((allys.sy2-n*ybar*ybar)/(n-1)); //sample standard deviation
#ifndef SILENT
   cout << "ybar,shat: " << ybar << ", " << shat <<  endl;
#endif

   //read x
   //the n*p numbers for x are stored as the p for first obs, then p for second, and so on.
	unsigned int myp=0, tempp=0;
	size_t p=0;
	std::vector<double> x;
	if(myrank!=0)
	{
		std::stringstream xfss;
		std::string xfs;
//		xfss << "x" << myrank << ".txt";
		xfss << xroot.str() << myrank << ".txt";
		xfs=xfss.str();
#ifndef SILENT
		cout << myrank << ": loading x's from " << xfs << endl;
#endif
		std::ifstream xf(xfs.c_str());  //file to read x from
   	while(xf >> xtemp) 
   	   x.push_back(xtemp);
   	myp = (unsigned int)x.size()/nn;
		p = (size_t) myp;
		if(x.size() != nn*p) {
      	cout << "error: input x file has wrong number of values\n"; 
			MPI_Finalize();
      	return 1;
		}
	}
	MPI_Reduce(&myp,&tempp,1,MPI_UNSIGNED,MPI_SUM,0,MPI_COMM_WORLD);
	if(myrank==0)
	{
		p=(size_t) tempp/(processes-1);
		if(p!=tempp/(processes-1))
		{
			cout << "error: input x files on compute nodes have different number of predictors p\n";
			MPI_Finalize();
			return 1;
		}
	}
#ifndef SILENT
	cout << myrank << ": myp=" << myp << " p=" << p << endl;
#endif

   //x for predictions
   dinfo dip; //data information for prediction
   dip.n=0;
   //Currently no out-of-sample prediction in MPIBART
#ifndef SILENT
   cout << "\nx for prediction read in:\n";
   cout << "n: " << dip.n << endl;
#endif
   if(dip.n) {
#ifndef SILENT
      cout << "first row: " <<  dip.x[0] << " ...  " << dip.x[p-1] << endl;
      cout << "last row: " << dip.x[(dip.n-1)*p] << " ...  " << dip.x[dip.n*p-1] << endl;
#endif
   }

   //--------------------------------------------------
   //optionally read in additional arguments
   size_t burn = 1000; //number of mcmc iterations called burn-in
   size_t nd = 1000; //number of mcmc iterations
   size_t thin = 10; //save a draw ever thin iterations after burn-in.
   size_t m=200; //number of trees in BART sum
   double lambda = 1.0; //this one really needs to be set
   double nu = 3.0;
   double kfac=2.0;
   if(argc>3) thin = atoi(argv[3]);
   if(argc>4) nd = atoi(argv[4]);
   if(argc>5) lambda = atof(argv[5]);
   if(argc>6) burn = atoi(argv[6]);
   if(argc>7) m = atoi(argv[7]);
   if(argc>8) nu = atof(argv[8]);
   if(argc>9) kfac = atof(argv[9]);

#ifndef SILENT
   cout <<"\nburn,nd,number of trees: " << burn << ", " << nd << ", " << m << endl;
   cout <<"\nlambda,nu,kfac: " << lambda << ", " << nu << ", " << kfac << endl;
#endif

   //--------------------------------------------------
   //x cutpoints
	std::vector<double> minx(p,INFINITY);
   std::vector<double> maxx(p,-INFINITY);
	xinfo xi;
	size_t nc=100; //100 equally spaced cutpoints from min to max.

	if(myrank!=0)
		makeminmax(p,nn,&x[0],minx,maxx);
	for(size_t i=0;i<p;i++) {		
		MPI_Allreduce(&minx[i],&minx[i],1,MPI_DOUBLE,MPI_MIN,MPI_COMM_WORLD);
		MPI_Allreduce(&maxx[i],&maxx[i],1,MPI_DOUBLE,MPI_MAX,MPI_COMM_WORLD);
#ifndef SILENT
		cout << myrank << ": minx[" << i << "]=" << minx[i] << " maxx[" << i << "]=" << maxx[i] << endl;
#endif
	}
	makexinfominmax(p,xi,nc,minx,maxx); //need to consruct the cutpoints on master and slaves.

   //--------------------------------------------------
   //trees
#ifndef SILENT
	cout << myrank << ": ybar=" << ybar << "\n";
#endif
   std::vector<tree> t(m);
   for(size_t i=0;i<m;i++) t[i].setm(ybar/m); //if you sum the fit over the trees you get the fit.

   //--------------------------------------------------
  //prior and mcmc
   pinfo pi;
   pi.pbd=1.0; //prob of birth/death move
   pi.pb=.5; //prob of birth given  birth/death

   pi.alpha=.95; //.95 prior prob a bot node splits is alpha/(1+d)^beta, d is depth of node
   pi.beta=2; //2 for bart means it is harder to build big trees.
   pi.tau=(maxy-miny)/(2*kfac*sqrt((double)m));
   pi.sigma=shat;

#ifndef SILENT
   cout << "\nalpha, beta: " << pi.alpha << ", " << pi.beta << endl;
   cout << "sigma, tau: " << pi.sigma << ", " << pi.tau << endl;
#endif

   //--------------------------------------------------
   //dinfo
	dinfo di;
	double *ftemp=NULL;
	double *r=NULL;
	double *allfit=NULL;
	
	if(myrank==0)
	{
		di.n=n; di.p=p; di.x=NULL; di.y=NULL;
	}
	else
	{
	   allfit = new double[nn]; //sum of fit of all trees
   	for(size_t i=0;i<nn;i++) allfit[i]=ybar;
	   r = new double[nn]; //y-(allfit-ftemp) = y-allfit+ftemp
	   ftemp = new double[nn]; //fit of current tree
	   di.n=nn; di.p=p; di.x = &x[0]; di.y=r; //the y for each draw will be the residual 		
	}

   //--------------------------------------------------
   //storage for ouput
   //in sample fit
	double *pmean;
	//for sigma draw
	double rss; //residual sum of squares
	double restemp; //a residual
  	//out of sample fit
 	double* ppredmean=0; //posterior mean for prediction
	double* fpredtemp=0; //temporary fit vector to compute prediction

	std::ofstream bsdf("bart-sd.txt"); //note that we write all burn+nd draws to this file.

	if(myrank!=0)
	{
	   pmean = new double[nn]; //posterior mean of in-sample fit, sum draws,then divide
   	for(size_t i=0;i<nn;i++) pmean[i]=0.0;

   	if(dip.n) {
      	ppredmean = new double[dip.n];
      	fpredtemp = new double[dip.n];
      	for(size_t i=0;i<dip.n;i++) ppredmean[i]=0.0;
   	}
	}

   //--------------------------------------------------
   //mcmc
	double temprss;
	MPI_Request request;
	MPI_Status status;
	if(myrank==0)
	{
		//Setup saving of the fitted trees
		std::ofstream treef("trees.txt");
		treef << xi << endl; //cutpoints
		treef << m << endl;  //number of trees
		treef << p << endl;  //dimension of x's
		treef << (int)(nd/thin) << endl; //number of draws from posterior that will be saved

#ifndef SILENT
		cout << "\nMCMC:\n";
#endif
		time_t tp;
		int time1 = time(&tp);
		for(size_t i=0;i<(nd+burn);i++)
		{
#ifndef SILENT
			if(i%100==0) cout << "i: " << i << endl;
#endif
			//draw trees
			for(size_t j=0;j<m;j++) 
			{
				bd(t[j],xi,pi,gen,processes-1);
				MPImasterdrmu(t[j],xi,pi,gen,processes-1);
			}
			//draw sigma
			rss=0.0;
			for(int j=1;j<=(processes-1);j++)
			{
				MPI_Recv(&temprss,1,MPI_DOUBLE,MPI_ANY_SOURCE,0,MPI_COMM_WORLD,&status);
				rss+=temprss;
			}
			pi.sigma = sqrt((nu*lambda + rss)/gen.chi_square(nu+n));
			if(i>=burn && !(((int)i)%((int)thin)))
			{
				bsdf << pi.sigma << endl;
				for(size_t j=0;j<m;j++) treef << t[j] << endl;  //save all the trees from this draw
			}
		} // end of mcmc loop

		int time2 = time(&tp);
#ifndef SILENT
		cout << "time for loop: " << time2-time1 << endl;
#endif
 	}
	else
	{
#ifndef SILENT
		cout << "slave node " << myrank << " is up" << endl;
#endif
		for(size_t i=0;i<(nd+burn);i++)
		{
			for(size_t j=0;j<m;j++)
			{
				fit(t[j],xi,di,ftemp);
				for(size_t k=0;k<nn;k++)
				{
					allfit[k] = allfit[k]-ftemp[k];
					r[k] = y[k]-allfit[k];
				}
				// Slave code to handle the bd() call on the master
				MPIslavegetsuff(t[j],xi,di);
				MPIslaveupdatebirthdeath(t[j]);
				// Slave code to handle the drmu() call on the master
				MPIslavedrmu(t[j],xi,di);
				fit(t[j],xi,di,ftemp);
				for(size_t k=0;k<nn;k++) allfit[k] += ftemp[k];
			}
			// Slave code to handle the draw sigma code on the master
			rss=0.0;
			for(size_t k=0;k<nn;k++) {restemp=y[k]-allfit[k]; rss += restemp*restemp;}
			MPI_Isend(&rss,1,MPI_DOUBLE,0,0,MPI_COMM_WORLD,&request);

			if(i>=burn)
			{
				for(size_t k=0;k<nn;k++) pmean[k] += allfit[k];
				if(dip.n)
				{
					for(size_t j=0;j<m;j++)
					{
						fit(t[j],xi,dip,fpredtemp);
						for(size_t k=0;k<dip.n;k++) ppredmean[k] += fpredtemp[k];
					}
				}
			}
			MPI_Wait(&request,MPI_STATUS_IGNORE); //can't go further until the rss send completed.

		} //end of mcmc loop

		for(size_t k=0; k<nn; k++) pmean[k] /= (double)nd;
		std::stringstream ofss;
		std::string ofs;
		ofss << "bart-fit" << myrank << ".txt";
		ofs=ofss.str();
#ifndef SILENT
		cout << myrank << ": saving fit to " << ofs << endl;
#endif
		std::ofstream bfitf(ofs.c_str());
		for(size_t i=0; i<nn; i++) bfitf << pmean[i] << endl;

		//cleanup
		delete[] allfit;
		delete[] r;
		delete[] ftemp;
		delete[] pmean;
		delete[] ppredmean;
		delete[] fpredtemp;
	}

	MPI_Finalize();
	return 0;
}

