## code adapted from http://probabilitynotes.wordpress.com
## and from McShane & Wyner (2010)
## Written to be flexible for number of pcs (<=10), but needs to have a modelN.bug
## file for each N

library(rjags)
setwd("bayesmodel")

mcshane.data <- read.table(paste("rjags_data_",numpc,".txt", sep = ""), head = T)
heads <- c("y[]", "Lagy[,1]", "Lagy[,2]", "PC[,1]", "PC[,2]", "PC[,3]", "PC[,4]","PC[,5]", "PC[,6]", "PC[,7]", "PC[,8]", "PC[,9]", "PC[,10]")
names(mcshane.data) <- heads[1:(numpc+3)]

## fixed parameters to for the priors
nrows <- 149
beta0 <- c(0*1:(numpc+3))
precbeta <- structure(.Data=c(
                        0.001, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0.001, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0.001, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0.001, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0.001, 0, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0.001, 0, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0.001, 0, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0.001, 0, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0, 0.001, 0, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0.001, 0, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.001, 0, 0,
                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.001, 0,
                        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.001
                        ), .Dim=c(13,13))[1:(numpc+3),1:(numpc+3)]

y <- mcshane.data[,grep("^y", names(mcshane.data))]
Lagy <- as.matrix(mcshane.data[,grep("Lagy", names(mcshane.data))])
PC <- as.matrix(mcshane.data[,grep("PC", names(mcshane.data))])

mcshane.all <- list(y = y, Lagy = Lagy, PC = PC,
                    nrows = nrows, beta0 = beta0, precbeta = precbeta)

mcshane.inits <- list(precsigma2 = 1, beta=c(0*1:(numpc+3)))

# note modelN.bug must exist, where N=numpc
mcshane.jags <- jags.model(file = paste("model",numpc,".bug", sep = ""),
                           data = mcshane.all,
                           inits = mcshane.inits,
                           n.chain = 2)

# burn in
update(mcshane.jags, 5000)

# monitor
mcshane.out <- jags.samples(mcshane.jags, c("beta", "mu", "precsigma"),
                            n.iter = 20000, thin = 20)

names(mcshane.out)
##[1] "beta"      "mu"        "precsigma"

# make it readable to coda
mcshane.beta <- as.mcmc.list(mcshane.out$beta)
mcshane.mu <- as.mcmc.list(mcshane.out$mu)
mcshane.precsigma <- as.mcmc.list(mcshane.out$precsigma)

# map to arrays used by McShane&Wyner code

post_beta <- mcshane.beta[[1]]
post_mu <- mcshane.mu[[1]]
post_sigma <- 1/sqrt(mcshane.precsigma[[1]])

# now use MW code for backcasting

# Create backcast with epsilon error:
tmp <- colMeans(post_beta)
int <- tmp[1]
betas <- tmp[-1]
yhat <- matrix(NA, dim(x_out)[1], numsim)
for(j in 1:numsim){
  tmp_x_out <- x_out
  for(i in 1:dim(yhat)[1]){
    yhat[i,j] <- int + sum(betas * tmp_x_out[i,]) + rnorm(1,0,post_sigma[j])
    tmp_x_out[i+1,1] <- yhat[i,j]
    tmp_x_out[i+2,2] <- yhat[i,j]
  }
  if(j%%25==0){print(paste("Complete: ", j, "/", numsim, sep=""))}
}
yhat1 <- yhat

# Create backcast with beta uncertainty:
yhat <- matrix(NA, dim(x_out)[1], numsim)
for(j in 1:numsim){
  tmp_x_out <- x_out
  for(i in 1:dim(yhat)[1]){
    tmpbetas <- post_beta[j,]
    yhat[i,j] <- tmpbetas[1] + sum(tmpbetas[-1] * tmp_x_out[i,])
    tmp_x_out[i+1,1] <- yhat[i,j]
    tmp_x_out[i+2,2] <- yhat[i,j]
  }
  if(j%%25==0){print(paste("Complete: ", j, "/", numsim, sep=""))}
}
yhat2 <- yhat

# Create backcast with beta uncertainty and epsilon error:
yhat <- matrix(NA, dim(x_out)[1], numsim)
for(j in 1:numsim){
  tmp_x_out <- x_out
  for(i in 1:dim(yhat)[1]){
    tmpbetas <- post_beta[j,]
    yhat[i,j] <- tmpbetas[1] + sum(tmpbetas[-1] * tmp_x_out[i,]) + rnorm(1,0,post_sigma[j])
    tmp_x_out[i+1,1] <- yhat[i,j]
    tmp_x_out[i+2,2] <- yhat[i,j]
  }
  if(j%%25==0){print(paste("Complete: ", j, "/", numsim, sep=""))}
}
yhat3 <- yhat

# Save output
save(list=c("yhat1","yhat2","yhat3","in_idxD","out_idxD","in_idxI","out_idxI","numpc","instrument"),file=paste("bayes_recon_",numpc,".R",sep=""))

setwd("../")