# Generate dataset according to Friedman function
#--------------------------------------------------
if(1) { #setup and simulate in-sample (x,y) and construct out-of-sample (x)
set.seed(99)
n=5000 #train (in-sample) sample size
p=5 #number of x's
sigma=1 #sd for observational error

X=matrix(runif(p*n),ncol=p) #in-sample x
#out-of-sample (or prediction) x:
sp=10 #number of prediction points per dimension to generate the prediction grid
Xp=as.matrix(expand.grid(seq(0,1,length=sp),seq(0,1,length=sp),seq(0,1,length=sp),seq(0,1,length=sp),seq(0,1,length=sp)))

f=10*sin(2*pi*X[,1]*X[,2])+(X[,3]-.5)^2+X[,4]+X[,5]
y=f+rnorm(n,sd=sigma)
fp=10*sin(2*pi*Xp[,1]*Xp[,2])+(Xp[,3]-.5)^2+Xp[,4]+Xp[,5]

npred=length(fp)
numnodes=8

pv=rep(0,numnodes-1) #number of in-sample observations for each slave node
pv[1:(numnodes-2)] = floor(n/(numnodes-1)) #number of slave = numnodes-1
pv[numnodes-1]=n-sum(pv)
cat("Generating ",numnodes-1," datasets of sizes ",pv,"\n")

# write out the dataset
for(i in 1:(numnodes-1)) #c indices of slaves
{
	if(i==1)
	{
		ii=1:pv[1] 
	} else {
		ii=(sum(pv[1:(i-1)])+1):sum(pv[1:i])
	}

	xx=X[ii,]
	yy=y[ii]
	write(yy,ncol=1,file=paste('y',i,'.txt',sep=""))
	write(t(xx),ncol=p,file=paste('x',i,'.txt',sep=""))
}

# write out the prediction points
pv=rep(0,numnodes) #number of out-of-sample observations for each slave node
pv[1:(numnodes-1)] = floor(nrow(Xp)/numnodes) #number of slaves = numnodes
pv[numnodes]=nrow(Xp)-sum(pv)
cat("Generating ",numnodes," predictive X datasets of sizes ",pv,"\n")

for(i in 1:numnodes) #note xpi.txt i=1,2,..numnodes
{
	if(i==1)
	{
		ii=1:pv[1]
	} else {
		ii=(sum(pv[1:(i-1)])+1):sum(pv[1:i])
	}
	
	xx=Xp[ii,]
	yy=fp[ii]
	write(t(xx),ncol=p,file=paste('xp',i,'.txt',sep=""))
}
}
#--------------------------------------------------
if(1) { # Fit Bart
nu=3
sighat = sd(y)

if(1) { #use least squares for sighat
df = data.frame(X,y)
lmf = lm(y~.,df)
sighat = summary(lmf)$sigma
}

sigq = .90
qchi = qchisq(1.0-sigq,nu)
lambda = (sighat*sighat*qchi)/nu #lambda parameter for sigma prior

nd=3500  #number of MCMC draws to "keep" (eg go into posterior means)
thin=100 #save draws ever "thin" MCMC draws after burn-in
burn=500 #number of initial draws called burn-in	
m= 200   #number of trees in bart sum
kfac=2   #bart prior parameter for determining prior stan dev of mu's

cmd = paste('mpirun -np',numnodes, ' ./mpibart x y ',thin,nd,lambda,burn,m,nu,kfac)
cat('cmd:\n')
cat(cmd)
cat('\n')
system(cmd)
}
#--------------------------------------------------
if(1) { # Plot in-sample fits
par(mfrow=c(1,1))

bfit=c()
for(i in 1:(numnodes-1)) {
fnm = paste('bart-fit',i,'.txt',sep='')
bfit = c(bfit,scan(fnm))
}
plot(f,bfit)
abline(0,1)
title('f(x) vs fhat(x), in-sample')

dev.copy2pdf(file='in-sample.pdf',height=6,width=12)
}
#--------------------------------------------------
if(1){ # Do the out-of-sample predictions
cmd = paste('mpirun -np',numnodes, './bartpred xp p')
cat('cmd:\n')
cat(cmd)
cat('\n')
system(cmd)
}
#--------------------------------------------------
if(1) { # Plot out-of-sample predictions
par(mfrow=c(1,1))
bfit=c()
for(i in 1:numnodes) {
fnm = paste('yp',i-1,'.txt',sep='')
temp = scan(fnm)
print(length(temp))
bfit = c(bfit,scan(fnm))
}
plot(fp,bfit,xlab="true friedman",ylab="est. friedman")
#plot(fp,bfit-2*bsd,pch=20,col="grey",xlab="true friedman",ylab="est. friedman")
#points(fp,bfit+2*bsd,pch=20,col="grey")
#points(fp,bfit,pch=20)
#abline(0,1)
#title('out-of-sample, f(x) vs. fhat(x)')
#dev.copy2pdf(file='out-of-sample.pdf',height=10,width=10)
}
