## Setup
DataOrig <- read.delim("http://web.as.uky.edu/statistics/users/pbreheny/701/data/cd4.txt")
J <- max(DataOrig$ID)
Data <- data.frame(Visit=rep(c(1, 4, 7, 10, 13, 16, 19), J), ID=rep(1:J, rep(7, J)))
Data$CD4Pct <- DataOrig$CD4Pct[match(apply(Data, 1, paste, collapse="-"), apply(DataOrig[,c("Visit", "ID")], 1, paste, collapse="-"))]
Data$y <- log(0.1+Data$CD4Pct)
n <- nrow(Data)
gData <- DataOrig[!duplicated(DataOrig$ID),]

## Varying-intercept, varying-slope model with ignorable missingness
jData <- list(y=Data$y, x=Data$Visit, Trt=gData$Treatment-1, n=n, J=J, ID=Data$ID)
model <- function() {
  ## Visit level
  for (i in 1:n) {
    y[i] ~ dnorm(a[ID[i]] + b[ID[i]]*x[i], sigma[1]^(-2))
  }
  
  ## Individual level
  for (j in 1:J) {
    a[j] ~ dnorm(ga[1] + ga[2]*Trt[j], sigma[2]^(-2))
    b[j] ~ dnorm(gb[1] + gb[2]*Trt[j], sigma[3]^(-2))
  }

  ## Priors
  for (j in 1:3) {
    sigma[j] ~ dunif(0,100)
  }
  for (j in 1:2) {
    ga[j] ~ dnorm(0, 0.0001)
    gb[j] ~ dnorm(0, 0.0001)
  }
  
  ## Monitor some y's
  y2[1:5] <- y[10:14]
}
require(R2jags); invisible(runif(1))
fit <- jags(model=model, param=c("ga", "gb", "sigma", "a", "b", "y2"), data=jData, n.iter=5000, n.thin=1)
attach.jags(fit)

## Convergence
max(gelman.diag(as.mcmc(fit))$psrf)
min(effectiveSize(as.mcmc(fit)))

## Posteriors
psm(ga)
psm(12*gb)
am1 <- apply(a, 2, mean)
bm1 <- apply(b, 2, mean)

## ID 2
with(subset(Data, ID==2), plot(Visit, y, pch=19, col="black", xlim=c(0, 19), ylab="log(CD4%)", ylim=c(-3, 2), las=1))
mtext(2)
abline(am1[2], bm1[2], col="slateblue", lwd=3)
points(Data$Visit[3:7], psm(y2)[,1], pch=19, col="gray")
for (i in 1:5) {
  lines(rep(Data$Visit[i+2],2), psm(y2)[i,2:3])
}
