## code to accompany the reading comprehension example from
## lecture 11 of PSTAT 215A, Bayesian Inference, UCSB
##
## Robert B. Gramacy

## load necessary libraries
library(mvtnorm)
source("rwish.R")

## read in the data file and extract sufficient stats
Y <- read.table("reading.txt", header=TRUE)
n <- nrow(Y)
ybar <- apply(Y, 2, mean)
Sdata <- cov(Y)
Sdatai <- solve(Sdata)
p <- ncol(Y)

## allocate memory for Gibbs Samples
S <- 5000
mu <- matrix(NA, nrow=S, ncol=2)
Sigma <- array(NA, dim=c(2,2,S))

## initialization
mu[1,] <- ybar
Sigma[,,1] <- Sdata

## prior settings
mu.0 <- c(50, 50)
Lambda.0 <- rbind(c(625, 312.5), c(312.5, 625))
Lambda.0i <- solve(Lambda.0)
nu.0 <- p + 2
S.0 <- Lambda.0

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

  ## update mu
  Sigmai <- solve(Sigma[,,s-1]) 
  Lambda.n <- solve( Lambda.0i + n*Sigmai )
  mu.n <- Lambda.n %*% ( Lambda.0i %*% mu.0 + n*Sigmai %*% ybar )
  mu[s, ] <- rmvnorm(1, mu.n, Lambda.n)

  ## update Sigma
  Ymmu <- t(Y) - mu[s,]
  S.n <- S.0 + Ymmu %*% t(Ymmu)
  Sigma[,,s] <- solve(rwish(nu.0 + n, solve(S.n)))
  
}


## checking for the posterior probability that the
## instruction was helpful

## plot the samples from the posterior for mu
plot(mu[,1], mu[,2], bty="n", main="")
abline(0,1, col=2, lty=2, lwd=2)
legend("bottomright", "no effect", col=2, lwd=2, lty=2, bty="n")

## posterior mean summary statistics
quantile(mu[,2] - mu[,1], prob=c(0.025, 0.975))
mean(mu[,2] > mu[,1])


## samples from the posterior predictive distribution
y <- matrix(NA, S, 2)
for(s in 1:S) {
  y[s,] <- rmvnorm(1, mu[2,], Sigma[,,s])
}

## plot the samples from the posterior predictive
plot(y[,1], y[,2], bty="n", main="")
abline(0,1, col=2, lty=2, lwd=2)
points(Y, col=2, pch=19)
legend("bottomright", "no effect", col=2, lwd=2, lty=2, bty="n")

## posterior predictive summary statistics
quantile(y[,2] - y[,1], prob=c(0.025, 0.975))
mean(y[,2] > y[,1])