###############################################
#Posterior Simulation Example
#Schools Data
#REVISED VERSION (see announcements page) - 2/1/06
###############################################

#Set working directory
setwd("U:\\Teaching\\Stat694\\Computing\\")

#Read in the schools data
schools.data <- read.table("schools.dat",header=T)
attach(schools.data)
N <- dim(schools.data)[1]    #Number of schools

###################################################################
#Monte Carlo (MC) Simulation, tau.sq known
###################################################################

#Function to sample from p(mu|tau.sq,y)
sample.mu <- function(tau.sq,y,N,sigma.sq){
                mu.post.mean <- sum((rep(1,N)/(sigma.sq+tau.sq))*y)/sum(rep(1,N)/(sigma.sq+tau.sq))
                mu.post.var <- 1/sum(rep(1,N)/(sigma.sq+tau.sq))         ####CORRECTION####
                return(rnorm(1,mean=mu.post.mean,sd=sqrt(mu.post.var)))
                }
#Function to sample from p(theta.j|mu,tau.sq,y)
sample.theta <- function(j,mu,tau.sq,y,sigma.sq){
                   theta.j.post.mean <- ((y[j]/sigma.sq[j])+(mu/tau.sq))/((1/sigma.sq[j])+(1/tau.sq))
                   theta.j.post.var <- ((1/sigma.sq[j])+(1/tau.sq))
                   return(rnorm(1,mean=theta.j.post.mean,sd=sqrt(theta.j.post.var)))
                   }
####
#Obtain samples from the posterior distriubion of the thetas and mu
####

#Set the value of tau.sq
tau.sq <- 100

#Set the number of samples and create variables to store the samples
K <- 1000
mu.post.samples <- rep(0,K)
theta.post.samples <- matrix(0,K,N)

#Draw samples from the posterior
for(k in 1:K){
      mu.post.samples[k] <- sample.mu(tau.sq,y,N,sigma.sq)
      for(n in 1:N){
            theta.post.samples[k,n] <- sample.theta(n,mu.post.samples[k],tau.sq,y,sigma.sq)
            }
      }

#Histogram of the samples from the posterior of mu
hist(mu.post.samples,main="Posterior of mu",xlab="mu")

#Boxplots of the posterior samples for theta and mu
boxplot(as.data.frame(cbind(theta.post.samples,mu.post.samples)),names=c(as.character(school),"mu"),main="Posterior of the thetas and the mu",ylim=range(y))
points(1:8,y,pch=19,col='green')

###################################################################
#Markov Chain Monte Carlo (MCMC) Simulation, tau.sq unknown
###################################################################

#Function to sample from p(tau.sq|mu,y)
sample.tau.sq <- function(mu,theta,N,alpha,beta){
                    alpha.post <- (N/2)+alpha
                    beta.post <- sum(.5*sum((theta-mu)^2)+beta)
                    return(1/rgamma(1,shape=alpha.post,scale=1/beta.post))
                    }

####
#Obtain samples from the posterior distriubion of the thetas, mu, tau
####

#Set hyperprior parameters
alpha <- .01
beta <- .01

#Set the number of samples and create variables to store the samples
K <- 1000
mu.post.samples <- rep(0,K)
theta.post.samples <- matrix(0,K,N)
tau.sq.post.samples <- rep(0,K)

#Set the initial values of the Markov Chain
mu.init <- 0
theta.init <- rep(0,N)
tau.sq.init <- 1

mu.post <- mu.init
theta.post <- theta.init
tau.sq.post <- tau.sq.init

#Draw samples from the posterior
for(k in 1:K){
     mu.post <- sample.mu(tau.sq.post,y,N,sigma.sq)
     mu.post.samples[k] <- mu.post
     for(n in 1:N){
            theta.post[n] <- sample.theta(n,mu.post,tau.sq.post,y,sigma.sq)
            }
     theta.post.samples[k,] <- theta.post
     tau.sq.post <- sample.tau.sq(mu.post,theta.post,N,alpha,beta)
     tau.sq.post.samples[k] <- tau.sq.post
     }

#Look at trace plots of the Markov chains
par(mfrow=c(10,1),mar=c(2,1,1,1))
plot(1:K,mu.post.samples,type='l',xlab="iter",ylab="mu")
for(i in 1:N){
      plot(1:K,theta.post.samples[,n],type='l',xlab="iter",ylab=paste("theta ",i))
      }
plot(1:K,tau.sq.post.samples,type='l',xlab="iter",ylab="tau")

par(mfrow=c(1,1),mar=c(1,1,1,1))
plot(mu.post.samples[1:10],type='l',xlab='iter',ylab='mu')

#Set the number of chains to keep
samples.keep <- 200:K

#Boxplots of the posterior samples for theta and mu
par(mfrow=c(1,1),mar=c(2,2,2,2))
boxplot(as.data.frame(cbind(theta.post.samples[samples.keep,],mu.post.samples[samples.keep])),names=c(as.character(school),"mu"),main="Posterior of the thetas and the mu",ylim=range(y))
points(1:8,y,pch=19,col='green')

#Histogram of tau.sq
hist(tau.sq.post.samples[samples.keep],xlab="tau.sq",main="Posterior distribution of tau.sq")

#Posterior median of tau.sq
quantile(tau.sq.post.samples[samples.keep],.5)