# Code written by Eric Ward, eric.ward@noaa.gov on 07/01/08
# This script uses either OpenBUGS or WinBUGS to do estimation of 
# the state-space multiple linear regression model with constant 
# observation & process error variance.
library("R2WinBUGS")

# covariates consist of predictors X (n x p, including 1s)
# response is univariate Y (n x 1)
# 
n = 10
p = 3
# Create simulated data for the example
X <- matrix(cbind(rep(1,n),seq(1,n),runif(n)),nrow=n,ncol=p)
Y <- rnorm(n= n, mean = X%*%c(0.7,0.3,0.1),1)
# Assign the prior means/variances to the 
B.prior_means <- rep(0,p)
B.prior_vars <- rep(10,p)
# Assign the shape / scale parameters
# Note: I'm using the ScaledInvChi2 distribution, following (Gelman 2004).
# Shape (v) interpreted as the amount of information the prior contributes,
# Scale (theta) interpreted as prior variance
# sig2 ~ ScaledInvChi2 (v, theta) equivalent to sig2 ~ InvGamma(v/2,theta/2)
# or tau = (1/sig2) ~ Gamma(shape = v/2, rate= v/2s2) with E[X] = 1/s2
# Observation error parameters:
sigmaO.v = 1
sigmaO.theta = 1
tauO.prior = c(sigmaO.v/2, (sigmaO.v/2)*sigmaO.theta)
# Process error parameters:
sigmaP.v = 1
sigmaP.theta = 1
tauP.prior = c(sigmaP.v/2, (sigmaP.v/2)*sigmaP.theta)

# this is the model string, generalized for ANY dataset
modelString = ("
model {
  for(i in 1:p) {
     B.tau[i] <- 1/B.prior_vars[i]
  }
  # prior on each precision parameter
  tauO ~ dgamma(tauO.prior[1],tauO.prior[2]);
  tauP ~ dgamma(tauP.prior[1],tauP.prior[2]);
  sigmaO <- sqrt(1/tauO);
  sigmaP <- sqrt(1/tauP);
  # process model
  for(i in 1:p) {B[1,i] ~ dnorm(B.prior_means[i],B.tau[i]); }
  for(i in 2:n) {
    # error in the regression parameters
    for(j in 1:p) {B[i,j] ~ dnorm(B[i-1,j],tauP); }
  } 
  # observation model
  for(i in 1:n) {
    pred[i] <- inprod(X[i,1:p],B[i,1:p]);
    Y[i] ~ dnorm(pred[i],tauO);
  }   
}
")

burn = 10000
chainLength = 10000
write(modelString,file="model.txt",append=FALSE)
data <- list("Y","X","n","p","B.prior_means","B.prior_vars","tauO.prior","tauP.prior");
parameters <- c("B","sigmaO","sigmaP","pred")
mod <- NA
try(mod <- bugs (data,inits=NULL,parameters,DIC=TRUE,"model.txt",n.chains=1,n.burnin=burn, 
   n.thin=1,n.iter=(burn+chainLength),program="openbugs",debug=FALSE),silent=TRUE)
attach.bugs(mod,overwrite=TRUE)

# look at model summary
mod

# plot the predicted trajectory
predicted <- matrix(0,n,3)
for(i in 1:n) predicted[i,] <- quantile(pred[,i],c(0.025,0.5,0.975))
plot(predicted[,2],main="Predicted State Vector w/95% Intervals",xlab="",ylab="",ylim=c(min(predicted,Y),max(predicted,Y)),type="l",lwd=2)
points(Y,col="red")
lines(predicted[,1],lty=3)
lines(predicted[,3],lty=3)