## Setup
Data <- read.delim("http://web.as.uky.edu/statistics/users/pbreheny/701/data/birthweight.txt")
attach(Data)
n <- nrow(Data)

## Descriptive
prop.table(table(Smk, THM),1)
prop.table(table(Eth, THM),1)
prop.table(table(Eth, Smk),1)

## Multivariate probit imputation model (takes quite a while)
jData <- list("lbw", "THM", "Male", "Dep", "Smk", "Eth", "SmkArea", "EthArea", "n")
model <- function() {
  for (i in 1:n) {
    lbw[i] ~ dbern(p[i])
    logit(p[i]) <- b[1] + b[2]*THM[i] + b[3]*Male[i] + b[4]*Dep[i] + b[5]*Smk[i] + b[6]*Eth[i]
  }
  for (k in 1:6) {
    b[k] ~ dnorm(0, 0.0001)
  }

  # Imputation
  for (i in 1:n) {
    Z[i,1:2] ~ dmnorm(mu[i,1:2], Omega[1:2,1:2])
    mu[i,1] <- d[1,1] + d[2,1]*THM[i] + d[3,1]*Male[i] + d[4,1]*Dep[i] + d[5,1]*SmkArea[i] + d[6,1]*EthArea[i]
    mu[i,2] <- d[1,2] + d[2,2]*THM[i] + d[3,2]*Male[i] + d[4,2]*Dep[i] + d[5,2]*SmkArea[i] + d[6,2]*EthArea[i]
    Smk[i] ~ dinterval(Z[i,1],0)
    Eth[i] ~ dinterval(Z[i,2],0)
  }
  Sigma[1,1] <- 1
  Sigma[2,2] <- 1
  Sigma[1,2] <- rho
  Sigma[2,1] <- rho
  rho ~ dunif(-1, 1)
  Omega <- inverse(Sigma)
  for (k in 1:6) {
    d[k,1] ~ dnorm(0, 0.0001)
    d[k,2] ~ dnorm(0, 0.0001)
  }
  Smk14 <- Smk[14]
  Eth14 <- Eth[14]
  Smk1476 <- Smk[1476]
  Eth1476 <- Eth[1476]
}
zinit <- function(x) {
  z <- rep(0, length(x))
  z[x == 1] <- 1
  z[x == 0] <- -1
  z
}
init <- function(){list(Z=cbind(zinit(Smk), zinit(Eth)))}
require(R2jags); invisible(runif(1))
fit <- jags(model=model, init=init, param=c("b", "d", "rho", "Smk14", "Eth14", "Smk1476", "Eth1476"), data=jData, n.iter=10000, n.thin=1)
attach.jags(fit)

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

## Posteriors
psm(exp(b), labels=c("Intercept", "THM", "Male", "Dep", "Smk", "Eth"))
psm(d[,,1], labels=c("Intercept", "THM", "Male", "Dep", "SmkArea", "EthArea"))
psm(d[,,2], labels=c("Intercept", "THM", "Male", "Dep", "SmkArea", "EthArea"))

## Birthweight model
B <- psm(exp(b))[-1,]
rownames(B) <- c("THM", "Male", "Dep", "Smk", "Eth")
CIplot(B, xlab="Odds ratio")
abline(v=1, col="gray")

## Missing-data models
B1 <- psm(d[,-1,1])
B2 <- psm(d[,-1,2])
rownames(B1) <- rownames(B2) <- c("THM", "Male", "Dep", "SmkArea", "EthArea")
CIplot(B1, sort=FALSE, xlab=~delta)
abline(v=0, col="gray")
CIplot(B2, sort=FALSE, xlab=~delta)
abline(v=0, col="gray")

## Imputation
psm(rho)
data.frame(Data[c(14, 1476),], SmkImp=c(mean(Smk14), mean(Smk1476)), EthImp=c(mean(Eth14), mean(Eth1476)))