# do lalonde example with randomForests and BART
library(BayesTree)

#--------------------------------------------------
#read in data
if(1) {
datadir = './' # change this to whatever director the data files are in
tr = read.table(paste(datadir,'nswre74_treated.txt',sep=''),header=F) #treated from randomized
psid = read.table(paste(datadir,'psid_controls.txt',sep=''),header=F) # controls from psid
nms = c('trt','age','edu','black','hisp','married','nodegree','re74','re75','re78') #variable names
names(tr)=nms
names(psid)=nms
} 

#--------------------------------------------------
#set up bart data 
if(1) { 
dat = rbind(tr,psid)
y = as.vector(dat[,10]) 
xtrain = as.matrix(dat[,1:9]) #x (inluding the treatment in first column)
#set up prediction x for treatment effect on the treated
x1 = as.matrix(tr[,1:9]); x1[,1]=1
x0 = as.matrix(tr[,1:9]); x0[,1]=0
xtest = rbind(x1,x0)
}

#-------------------------------------------------
# run bart 
if(1) {  
   set.seed(99)
   if(0) { # if 0 do a short run, else longer, long run is what is reported in the paper
      tbart = system.time({
      rbart = bart(xtrain,y,xtest,nskip=2000,ndpost=10000,keepevery=10)
      })
   } else {
      rbart = bart(xtrain,y,xtest)
   }
   if(1) { # plot sigma draws to check convergence of mcmc
      par(mfrow=c(1,1))
      plot(c(rbart$first.sigma,rbart$sigma))
   }
}

#--------------------------------------------------
# get bart results
if(1) {
nd = nrow(rbart$yhat.test)
nx = ncol(rbart$yhat.test)/2
diff = rbart$yhat.test[,1:nx] - rbart$yhat.test[,1:nx + nx]
pTr = apply(diff,1,mean) # draws from posterior of average treatment effect on the treated
eTr = apply(diff,2,mean) # posterior mean of treatment effect for x of each treated subject
}

#--------------------------------------------------
# plot posterior of average effect on treated
if(1) {
par(mfrow=c(1,1))
hist(pTr,xlab='average effect on treated',main='')
}