## code to accompany the Normal midge example from 
## lecture 7-8 of PSTAT 215A, Bayesian Inference, UCSB
##
## Robert B. Gramacy adapted from Peter Hoff

## prior parameters
mu.0 <- 1.9
tau2.0 <- 0.95
tau2.0i <- 1/tau2.0

## data
y <- c(1.64, 1.70, 1.72, 1.74, 1.82, 1.82, 1.82, 1.90, 2.08)

## calculate the data sufficient statistics
y.bar <- mean(y)
n <- length(y)
s2 <- var(y)

## posterior statistics
mu.n <- (tau2.0i * mu.0 + (n/s2)*y.bar)/(tau2.0i + n/s2)
tau2.n <- 1/(tau2.0i + n/s2)

## plot the prior and posterior densities
mu <- seq(0,3.5,length=1000)
mu.prior <- dnorm(mu, mean=mu.0, sd=sqrt(tau2.0))
mu.post <- dnorm(mu, mean=mu.n, sd=sqrt(tau2.n))
plot(mu, mu.post, type="l", lwd=2, main="",
     xlab="mu", ylab="density", bty="n")
lines(mu, mu.prior, col=2, lty=2, lwd=2)
legend("right", c("posterior" , "prior"), col=1:2,
       lty=1:2, lwd=2, bty="n")

## now joint posterior

## setting up the prior
k.0 <- 1
s2.0 <- 0.01; nu.0 <- 1

## posterior sufficient statistics
k.n <- k.0 + n
nu.n <- nu.0 + n
s2.n <- nu.0*s2.0 + (n-1)*s2 + k.0*n*(y.bar - mu.0)^2/k.n
s2.n <- s2.n / nu.n

## make a function for calculating the posterior probability
post1 <- function(theta, mu.n, k.n, s2.n, nu.n)
  {
    p <- dnorm(theta$mu, mean=mu.n, sd=sqrt(theta$s2/k.n))
    p <- p * dgamma(1/theta$s2, nu.n/2, nu.n*s2.n/2)
  }

## make the posterior image plot
library(akima)
N <- 200
mu.grid1 <- seq(1.65, 1.95, length=N)
s2.grid1 <- seq(0.006, 0.06, length=N)
g1 <- expand.grid(mu.grid1, s2.grid1)
names(g1) <- c("mu", "s2")
p1 <- matrix(post1(g1, mu.n, k.n, s2.n, nu.n), nrow=N)
image(mu.grid1, s2.grid1, p1)

## overlaying samples from the joint posterior
S <- 10000
s2.mc <- 1/rgamma(S, nu.n/2, s2.n*nu.n/2)
mu.mc <- rnorm(S, mu.n, sqrt(s2.mc/k.n))
points(mu.mc[1:1000], s2.mc[1:1000], pch=19, cex=0.5)
abline(v=mean(mu.mc))
abline(h=mean(s2.mc))

## plotting marginal posterior (kenrel) densities
plot(density(s2.mc), xlim=c(0.006, 0.06), bty="n",
     lwd=2, main="", xlab="s2", ylab="density")
plot(density(mu.mc), xlim=c(1.65, 1.95), bty="n",
     lwd=2, main="", xlab="mu", ylab="density")

## exploring posterior credible intervals
qmu.cred <- quantile(mu.mc, c(0.025, 0.975))
qmu.conf <- qt(c(0.025, 0.975), df=n-1)*(sqrt(s2)/sqrt(n)) + y.bar
abline(v=qmu.cred, lwd=2, lty=2, col=2)
abline(v=qmu.conf, lty=3, col=3, lwd=2)
legend("topright", c("cred", "conf"), lty=2:3, col=2:3,
       lwd=2, bty="n")


## Gibbs Sampling

## Initialization
S <- 1000
mu.gs <- s2.gs <- rep(NA, S)
mu.gs[1] <- y.bar
s2.gs[1] <- s2
nu.n <- nu.0 + n

## Gibbs sampling
for(s in 2:S) {

  ## calculate mean and variance parameters of mu (Normal) conditional
  tau2.n <- 1/(1/tau2.0 + n/s2.gs[s-1])
  mu.n <- (mu.0/tau2.0 + n*y.bar/s2.gs[s-1]) * tau2.n

  ## generate a new mu value
  mu.gs[s] <- rnorm(1, mu.n, sqrt(tau2.n))

  ## calcuate IG parameters for s2 conditional
  s2.n <- (nu.0*s2.0 + (n-1)*s2 + n*(y.bar - mu.gs[s])^2)/ nu.n

  ## draw new s2 value
  s2.gs[s] <- 1/rgamma(1, nu.n/2, nu.n*s2.n/2)
}


## post2:
##
## calculate the full posterior for the indepent prior

post2 <- function(theta, y, mu.0, tau2.0, nu.0, s2.0)
  {
    mu.p <- dnorm(theta$mu, mu.0, sqrt(tau2.0)) 
    s2.p <- dgamma(1/theta$s2, nu.0/2, s2.0*nu.0/2)
    lik <- rep(NA, nrow(theta))
    for(i in 1:nrow(theta)) {
      lik[i] <- prod(dnorm(y, theta$mu[i], sqrt(theta$s2[i])))
    }
    return(mu.p * s2.p * lik)
  }


## make plot of all samples with true joint density underneath
mu.grid2 <- seq(1.6, 2, length=N)
s2.grid2 <- seq(0.002, 0.15, length=N)
g2 <- expand.grid(mu.grid2, s2.grid2)
names(g2) <- c("mu", "s2")
p2 <- matrix(post2(g2, y, mu.0, tau2.0, nu.0, s2.0), nrow=N)
image(mu.grid2, s2.grid2, p2)
points(mu.gs, s2.gs)

## summarize the marginal means
mean(mu.gs)
mean(s2.gs)

## summarize quantiles from the gibbs samples
quantile(mu.gs, c(0.025, 0.975))
quantile(s2.gs, c(0.025, 0.975))