## 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; informative missingness
jData <- list(y=Data$y, x=Data$Visit, Trt=gData$Treatment-1, n=n, J=J, ID=Data$ID, Miss=is.na(Data$y))
model <- function() {
  ## Visit level
  for (i in 1:n) {
    y[i] ~ dnorm(a[ID[i]] + b[ID[i]]*x[i], sigma[1]^(-2))
    Miss[i] ~ dbern(p[i])
    logit(p[i]) <- gm[1] + gm[2]*y[i]
  }
  gm[1] ~ dnorm(0, 0.0001)
  gm[2] ~ dnorm(0, 10)
  
  ## 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", "gm"), 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)
am2 <- apply(a, 2, mean)
bm2 <- apply(b, 2, mean)

## Plot missing data relationship
psm(gm)
mm <- apply(gm, 2, mean)
xx <- seq(-3, 4, len=99)
plot(exp(xx), ilogit(mm[1]+mm[2]*xx), type="l", ylim=c(0,1), lwd=3, las=1, xlab="CD4 %", ylab="Pr(Miss)")

## ID 2 (am1, bm1 come from ignorable model)
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))
abline(am1[2], bm1[2], col=col[1], lwd=3)
abline(am2[2], bm2[2], col=col[2], 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])
}
toplegend(c("Ignorable", "Informative"), lwd=3, col=col, ncol=2)

## Lines
ind <- as.numeric(names(sort(table(DataOrig$ID))[c(1:4, 51:54, 101:104, 151:154, 201:204)]))
par(mfcol=c(4,5), mar=c(5, 5, 0, 1))
col <- pal(2)
for (j in ind) {
  with(subset(Data, ID==j), plot(Visit, y, pch=19, col="gray", xlim=c(0, 19), ylab="log(CD4%)", ylim=c(mean(y, na.rm=TRUE)+c(-1,1)*2), las=1))
  abline(am1[j], bm1[j], col=col[1], lwd=3)
  abline(am2[j], bm2[j], col=col[2], lwd=3)
}

## Loglines
par(mfcol=c(4,5), mar=c(5, 5, 0, 1))
ind <- as.numeric(names(sort(table(DataOrig$ID))[c(1:4, 51:54, 101:104, 151:154, 201:204)]))
col <- pal(2)
logline <- function(a, b, ...) {
  xx <- seq(0, 19, len=99)
  lines(xx, exp(a+b*xx), ...)
}
for (j in ind) {
  with(subset(Data, ID==j), plot(Visit, y, pch=19, col="gray", xlim=c(0, 19), ylab="CD4%", ylim=c(0, 50), las=1))
  logline(am1[j], bm1[j], col=col[1], lwd=3)
  logline(am2[j], bm2[j], col=col[2], lwd=3)
}