#--------------------------------------------------
if(1) { #friedman function, simulate data and write out files
f = function(x){
   10*sin(pi*x[,1]*x[,2]) + 20*(x[,3]-.5)^2+10*x[,4]+5*x[,5]
}

set.seed(14)
n = 100      #train sample size
np=100       #test sample size
p=10         #number of x variables
x=matrix(runif(n*p),n,p) 
xp=matrix(runif(np*p),np,p) 

fy = f(x)
fpy = f(xp)

sigma = 1.0  
y=fy+sigma*rnorm(n)
ybar=mean(y)
y=y-ybar #bart assumes y is centered
fy=fy-ybar
fpy = fpy-ybar

write(y,file='y.txt',ncol=1)
write(t(x),file='x.txt',ncol=p)
write(t(xp),file='xp.txt',ncol=p)
write(fy,file='fy.txt',ncol=1)
write(fpy,file='fpy.txt',ncol=1)

#run regression and get regression estimate of sigma (sighat)
df = data.frame(x,y)
lmf = lm(y~.,df)
sighat = summary(lmf)$sigma
}
#--------------------------------------------------
if(1) { #run bart
#system('make bart') #compile bart

nd=1000 #number of draws after burn-in
nu=3    #degrees of freedom parameter for sigma prior

#choose lambda using bart default prior setup
sigq = .9
qchi = qchisq(1.0-sigq,nu)
lambda = (sighat*sighat*qchi)/nu #lambda parameter for sigma prior

burn=100 #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('./bart x.txt y.txt xp.txt',nd,lambda,burn,m,nu,kfac)
cat('cmd:\n')
cat(cmd)
system(cmd)
}
#--------------------------------------------------
if(1) { #plot bart output
#set up plotting
par(mfrow=c(2,2))
par(mai=c(.8,.5,.5,.2))
par(omi=c(0,0,0,0))

#read-in bart output
sd = scan('bart-sd.txt') #draws of sigma
bf = scan('bart-fit.txt') #posterior mean of f(x), in-sample
bpf = scan('bart-pred-fit.txt') #posterior mean of f(x), out-of-sample

#plot sigma draws
plot(sd)
abline(h=sigma,col='blue')
title(main='sigma draws including burn-in',cex.main=.7)

#plot in-sample fit
plot(fy,bf)
abline(0,1,col='red')
title(main='f estimate vs. true, in-samp',cex.main=.7)

#plot out-of-sample fit
plot(fpy,bpf)
abline(0,1,col='red')
title(main='f estimate vs. true, out-samp',cex.main=.7)
}
