##############################################################################
#
# Multiple normal linear regression - Gibbs sampler
#
##############################################################################
set.seed(12345)
n          = 200
beta.true  = c(1.0,0.8,0.2)
sigma.true = 0.5
p = length(beta.true)
X = matrix(rnorm(n*p),n,p)
y = X%*%beta.true+rnorm(n,0,sigma.true)

# OLS/ML inference
# ----------------
fit.ols   = lm(y~X-1)
beta.ols  = fit.ols$coef
sigma.ols = summary(fit.ols)$sigma
c(beta.ols,sigma.ols)

# Sufficient statistics
iV0 = solve(V0)
iV0beta0 = iV0%*%beta0
XtX = t(X)%*%X
Xty = t(X)%*%y

# Bayesian inference via Gibbs sampler
# ------------------------------------

# Prior of beta   ~ N(beta0,V0) 
beta0  = rep(0,p)
V0     = diag(9,p)


# Prior of sigma2 ~ IG(alpha0,gamma0)
# E,DP are the mean and standard deviation
# of the inverse gamma distribution
# alpha0 and gamma0 are derived from them
E      = 0.5
DP     = 1.5
alpha0 = 2+E^2/DP^2
gamma0 = (alpha0-1)*E

# Drawing from full conditional distributions
draw.full.sigma2 = function(beta){
  ystar = y-X%*%beta
  alpha1 = alpha0+n/2
  gamma1 = gamma0+sum(ystar^2)/2
  return(1/rgamma(1,alpha1,gamma1))
}
draw.full.beta = function(sigma2){
  V1 = solve(iV0+XtX/sigma2)
  beta1 = V1%*%(iV0beta0+Xty/sigma2)
  return(beta1+t(chol(V1))%*%rnorm(p))
}

# Initial value for beta and Gibbs sampler set up
set.seed(2355)
burnin = 1000
M      = 1000
niter  = burnin+M
draws  = matrix(0,niter,p+1)
beta   = beta.ols
for (iter in 1:niter){
  sigma2 = draw.full.sigma2(beta)
  beta   = draw.full.beta(sigma2)
  draws[iter,] = c(beta,sqrt(sigma2))
}
draws1 = draws[(burnin+1):niter,]

beta.post = apply(draws1[,1:p],2,mean)
sigma.post = mean(draws1[,p+1])

# Trace plots
par(mfrow=c(3,p+1),mai=c(0.8,0.8,0.2,0.1))

for (i in 1:p){
  ts.plot(draws[,1],main=paste("beta ",i,sep=""),
          ylab="",xlab="Gibbs sampler iteration")
  abline(v=burnin,col=2,lwd=3,lty=2)
}
ts.plot(draws[,p+1],main="sigma",ylab="",xlab="Gibbs sampler iteration")
abline(v=burnin,col=2,lwd=3,lty=2)

for (i in 1:(p+1)){
  acf(draws1[,i],main="")
}

for (i in 1:p){
  hist(draws1[,i],prob=TRUE,main="",xlab="",ylab="Posterior density of beta")
  points(beta.true[i],0,col=2,pch=16,cex=1.2)
  points(beta.ols[i],0,col=3,pch=16,cex=1.2)
  points(beta.post[i],0,col=4,pch=16,cex=1.2)
}

hist(draws1[,p+1],prob=TRUE,main="",xlab="",ylab="Posterior density of sigma")
points(sigma.true,0,col=2,pch=16,cex=1.2)
points(sigma.ols,0,col=3,pch=16,cex=1.2)
points(sigma.post,0,col=4,pch=16,cex=1.2)
legend("topleft",legend=c("TRUE","MLE","BAYES"),col=2:4,bty="n",pch=16,cex=1.2)


#pairs(draws1)