/*
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>

// for rand()
#include "stdio.h"
#include "stdlib.h"
#include "time.h"

#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);

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

	if(argc<2) {cout << "You need to specify at least the number of Monte Carlo total integration points\n"; return 1;}

	//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];  //all the trees


	//Setup the random number generator seeds
	size_t N,n;
	uint seed;
	double rnum=0.0;
	std::vector<double> xp1,xp2;  //integration points, stored like x

	N=atoi(argv[1]); //number of Monte Carlo integration points over all the inputs.
	n=N/processes;   //number of integration points to be computed on each compute node.
//	seed=time(NULL);
//	seed=seed+myrank;
	seed=myrank;
	RNG gen(seed);  //initialize random number generator

	cout << myrank << ": N=" << N << " total integration points of dimension " << p << ", n=" << n << " on this node" << endl;
	cout << myrank << ": seed=" << seed << endl;

	time_t tp;
	int time1=time(&tp);	

	//Construct the integration points for this node (matrices X1 and X2 in Homma & Saltelli &/or Saltelli text)
	for(size_t i=0;i<n*p;i++)
	{
		rnum=gen.uniform();//-.5;
		xp1.push_back(rnum);  //we assume the x's are in [-.5,.5].
	}
	for(size_t i=0;i<n*p;i++)
	{
		rnum=gen.uniform();//-.5;
		xp2.push_back(rnum);
	}
	std::vector<double> xp1t(xp1);  //copies of
	std::vector<double> xp2t(xp2);  //the 2 prediction sets

	dinfo dip1, dip2;  //data information for prediction
	dip1.n=0;
	dip2.n=0;
	size_t np = xp1.size()/p;
	if(xp1.size() != np*p) { 
		cout << "error, wrong number of elements in prediction data set\n";
		return 1;
	}
	if(np)
		dip1.n=np; dip1.p=p; dip1.x=&xp1t[0]; dip1.y=0; //there are no y's!
		dip2.n=np; dip2.p=p; dip2.x=&xp2t[0]; dip2.y=0;

	//Compute OAT variance components D0,..,Dp and TOAT variance components D_ij and corresponding sensitivites
	//and the total sensitivity indices TS_i

	//Do this nodes predictions
	double* ppredmean1; //posterior mean for prediction
	double* ppredmean2;
	double f0=0.0,D=0.0,Di=0.0,Dni=0.0,allf0=0.0,f02=0.0,m0=0.0,m02=0.0,gp;
	size_t numgps=10; //number of grid points for the 1-way summaries
	if(argc==3) numgps=(size_t)atoi(argv[2]);
	ppredmean1 = new double[dip1.n];
	ppredmean2 = new double[dip2.n];

	//compute the p 1-way summaries at a grid of points -------------------------------------------------------
	std::stringstream ofss;
	std::string ofs;
	std::ofstream *saf;
	if(myrank==0) {
		ofss << "saoneway.txt";
		ofs=ofss.str();
		cout << myrank << ": saving 1-way sensitivities to " << ofs << endl;
		saf=new std::ofstream(ofs.c_str());
	}

	for(size_t i=0;i<p;i++)
	{
		if(myrank==0) cout << "Dimension " << i+1 << endl;
		for(size_t j=0;j<=numgps;j++)
		{
			xp2t=xp2;
			gp=(double)j/numgps;
			//gp=-.5+(double)j/numgps;
			for(size_t k=0;k<n;k++)
				xp2t[p*k+i]=gp;
			makepostpred(dip2,xi,t,&ppredmean2[0]);
			f0=0.0; f02=0.0; m0=0.0; m02=0.0;
			for(size_t k=0;k<n;k++)
			{
				f0=f0+ppredmean2[k];
				f02=f02+ppredmean2[k]*ppredmean2[k];
			}
			MPI_Reduce(&f0,&m0,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
			MPI_Reduce(&f02,&m02,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
			if(myrank==0) {
				m0=m0/N;
				m02=m02/N;
				m02=m02-m0*m0;
				*saf << i << " " << gp << " " << m0 << " " << sqrt(m02) << endl;
			}
		}
	}
	if(myrank==0) delete saf;
	// end of 1-way summaries --------------------------------------------------------------------------------


	//compute f0^2
	if(myrank==0) cout << "Calculating f0^2" << endl;
	f0=0.0; f02=0.0;
	makepostpred(dip1,xi,t,&ppredmean1[0]);
	for(size_t i=0;i<n;i++)
		f0=f0+ppredmean1[i];
	MPI_Reduce(&f0,&allf0,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
	if(myrank==0) {
		allf0=allf0/N;
		f02=allf0*allf0;
	}


	//compute D
	if(myrank==0) cout << "Calculating D" << endl;
//	xp2t=xp2;
//	makepred(dip2,xi,t,&ppredmean2[0]);
	f0=0.0;
	for(size_t i=0;i<n;i++)
		f0=f0+ppredmean1[i]*ppredmean1[i]; //ppredmean2[i];
	MPI_Reduce(&f0,&D,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
	if(myrank==0) {
		D=D/N-f02;
		cout << "D=" << D << endl;
	}


	//compute Di's and Si's
	if(myrank==0) cout << "Calculating Di's and Si's" << endl;
	for(size_t i=0;i<p;i++)
	{
		xp2t=xp2;
		for(size_t j=0;j<n;j++)
			xp2t[p*j+i]=xp1t[p*j+i];
		makepostpred(dip2,xi,t,&ppredmean2[0]);
		f0=0.0; Di=0.0;
		for(size_t j=0;j<n;j++)
			f0=f0+ppredmean1[j]*ppredmean2[j];
		MPI_Reduce(&f0,&Di,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
		if(myrank==0) {
			Di=Di/N-f02;
			cout << "D" << i << "=" << Di << endl;
			cout << "S" << i << "=" << Di/D << endl;
		}
	}


	//compute D~i's and TSi's
	if(myrank==0) cout << "Calculating D~i's and TSi's" << endl;
	for(size_t i=0;i<p;i++)
	{
		xp2t=xp1;
		for(size_t j=0;j<n;j++)
			xp2t[p*j+i]=xp2[p*j+i];
		makepostpred(dip2,xi,t,&ppredmean2[0]);
		f0=0; Dni=0;
		for(size_t j=0;j<n;j++)
			f0=f0+ppredmean1[j]*ppredmean2[j];
		MPI_Reduce(&f0,&Dni,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
		if(myrank==0)
		{
			Dni=Dni/N-f02;
			cout << "TS" << i << "=" << 1-Dni/D << endl;
		}
	}

	//We could write out the sensitivites in some data file here if we wanted to,
	//but currently we're just printing them to stdout as they are computed above.
	int time2=time(&tp);
	if(myrank==0)	cout << "time for sensitivity analysis: " << time2-time1 << endl;

	//cleanup
	delete[] ppredmean1;
	delete[] ppredmean2;

	MPI_Finalize();

	return 0;
} //end of main
