## Read in data
AllData <- read.delim("http://web.as.uky.edu/statistics/users/pbreheny/701/data/radon.txt")
Data <- subset(AllData, state=="MN")
county <- droplevels(Data$county)
x <- Data$floor
cid <- as.numeric(county) ## County ID
n <- nrow(Data)
J <- length(levels(county))

## Read in county-level soil uranium data
AllData2 <- read.delim("http://web.as.uky.edu/statistics/users/pbreheny/701/data/radon-county.txt")
Data2 <- subset(AllData2, st=="MN")
u <- log(Data2$Uppm[match(levels(county), Data2$cty)])

## Outcome
radon <- Data$activity
y <- log(ifelse(radon==0,.1, radon))

## Group level predictor
jData <- list("y", "x", "cid", "n", "J", "u")
model <- function() {
  ## Likelihood
  for (i in 1:n) {
    y[i] ~ dnorm(a[cid[i]] + b*x[i], sigma[1]^(-2))
  }
  ## Priors
  for (j in 1:J) {
    a[j] ~ dnorm(g[1] + g[2]*u[j], sigma[2]^(-2))
  }
  b ~ dnorm(0, 0.001)
  for (j in 1:2) {
    g[j] ~ dnorm(0, 0.001)
    sigma[j] ~ dunif(0, 100)
  }  
}
require(R2jags); invisible(runif(1))
fit <- jags(model=model, param=c("g", "sigma", "a", "b"), data=jData, n.iter=10000, n.thin=1)
plot(fit)
fit
attach.jags(fit)

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

## County-level model
u.rep <- u[cid]
fit1 <- lm(y~u.rep+x)
plot(u, apply(a, 2, median), ylim=c(0.5, 2.1), pch=19, cex=0.5, las=1, ylab=~alpha)
abline(median(g[,1]), median(g[,2]), col=pal(2)[2], lwd=2)
abline(coef(fit1)[1], coef(fit1)[2], col=pal(2)[1], lwd=2)
for (j in 1:85) {
  lines(rep(u[j], 2), quantile(a[,j], c(.1, .9)), col="gray60")
}
toplegend(c("OLS", "Hierarchical"), ncol=2, col=pal(2), lwd=2)

## Posterior for gamma[2]
quantile(g[,2], c(.5, .025, .975))
2^quantile(g[,2], c(.5, .025, .975))

## Variance components
apply(sigma, 2, mean)
icc <- sigma[,2]^2/(sigma[,1]^2+sigma[,2]^2)
mean(icc)

## Comparison with original HLM
## aa comes from 3-19 lecture
nn <- as.numeric(table(county))
plot(aa, apply(a, 2, median), pch=16, cex=2*log(nn)/max(log(nn)), xlab=~alpha*", original HLM", ylab=~alpha*", Uranium HLM")
abline(0,1,col="gray60")

## Not identifiable:
fit2 <- lm(y~0+u.rep+x+county)