##################################################################################################
#
#  Bayesian linear regression with Laplace-Gamma prior for (beta,sigma2)
#
#  Goal: Show a two-step MCMC algorithm
#
#       1) Sample beta|sigma2: independent MH algorithm, proposal = N(betahat,sigma2*inv(X'X))
#       2) Sample sigma2|beta: random-walk MH algorithm, proposal = Gamma(10,10/sigma2.old)  
#
#
# By Hedibert Freitas Lopes
# This version: May 28th 2021
#
##################################################################################################

library(mvtnorm)
 


################################################################
# DATA: WNBA/NBA 2013-2014 Player Heights, Weights and Positions
################################################################

nba     = read.csv("http://hedibert.org/wp-content/uploads/2016/03/height-weight-nba.csv")
wnba    = read.csv("http://hedibert.org/wp-content/uploads/2016/03/height-weight-wnba.csv")

n1 = nrow(wnba)
n2 = nrow(nba)
wnba.F = rep(0,n1)
wnba.C = rep(0,n1)
nba.F  = rep(0,n2)
nba.C  = rep(0,n2)
wnba.F[wnba[,3]=="F"]=1
wnba.C[wnba[,3]=="C"]=1
wnba.F[wnba[,3]=="GF"]=1
wnba.C[wnba[,3]=="FC"]=1
nba.F[nba[,2]=="F"]=1
nba.C[nba[,2]=="C"]=1

n = n1+n2
p = 8
y = c(wnba[,5],nba[,4])
y = (y-mean(y))/sd(y)
X = matrix(0,n,p)
X[,1] = 1
X[,2] = c(wnba[,4],nba[,3])
X[,2] = (X[,2]-mean(X[,2]))/sd(X[,2])
X[,3] = c(rep(0,n1),rep(1,n2))
X[,4] = c(wnba.F,nba.F)
X[,5] = c(wnba.C,nba.C)
X[,6] = X[,2]*X[,3]
X[,7] = X[,2]*X[,4]
X[,8] = X[,2]*X[,5]

namesX = c("Icpt","height","sex","forward","center","height*sex","height*forward","height*center")

pos = rep(1,n)
pos[X[,4]==1]=2
pos[X[,5]==1]=3

par(mfrow=c(1,2))
plot(X[,2],y,col=X[,3]+1,xlab="Height (standardized)",ylab="Weight (standardized)",lwd=2)
legend("topleft",legend=c("WNBA","NBA"),col=1:2,lwd=2,bty="n")
plot(X[,2],y,col=pos,xlab="Height (standardized)",ylab="Weight (standardized)",lwd=2)
legend("topleft",legend=c("Guard","Forward","Center"),col=1:3,lwd=2,bty="n")

################
# OLS ESTIMATION
################
summary(lm(y~X-1))

# OLS estimation
iXtX  = solve(t(X)%*%X)
bhat  = iXtX%*%t(X)%*%y
s2    = sum((y-X%*%bhat)^2)/(n-p)
var.b = s2*iXtX
root  = t(chol(var.b))


##########################################################
# BAYESIAN INFERENCE: IND.MH for beta and RW.MH for sigma2
##########################################################
log.likelihood = function(beta,sigma2){
  -0.5*n*log(sigma2)-0.5*sum((y-X%*%beta)^2)/sigma2
}

log.prior = function(beta,sigma2){
  -p*tau-(1/tau)*sum(abs(beta))*dgamma(sigma2,a,b)
}

prior = function(beta){
(1/(2*tau))*exp(-abs(beta)/tau)
}

# Prior hyperparameters
a   = 1
b   = 1
tau = 1

# Initial values
sigma2 = s2
beta   = as.vector(bhat)

# MCMC set up
M0      = 1000
M       = 1000
lag     = 10
niter   = M0+lag*M
betas   = matrix(0,niter,p)
sigma2s = rep(0,niter)

# MCMC algorithm
set.seed(12345)
for (iter in 1:niter){
  # Sampling beta
  beta.star = as.vector(bhat + root%*%rnorm(p))
  logacceptance = min(0,log.likelihood(beta.star,sigma2)+log.prior(beta.star,sigma2)-
                        dmvnorm(beta.star,bhat,var.b,log=TRUE)-
                        log.likelihood(beta,sigma2)-log.prior(beta,sigma2)+
                        dmvnorm(beta,bhat,var.b,log=TRUE))
  if (log(runif(1))<logacceptance){
    beta = beta.star	 
  }
  # Sampling sigma2
  sigma2.star = rgamma(1,10,10/sigma2)
  logacceptance = min(0,log.likelihood(beta,sigma2.star)+log.prior(beta,sigma2.star)-
  dgamma(sigma2.star,10,10/sigma2,log=TRUE)-
  log.likelihood(beta,sigma2)-log.prior(beta,sigma2)+dgamma(sigma2,10,10/sigma2.star,log=TRUE))
  if (log(runif(1))<logacceptance){
    sigma2 = sigma2.star
  }
  betas[iter,] = beta
  sigma2s[iter] = sigma2
}

# Trace plots
par(mfrow=c(3,3),mai=c(0.3,0.3,0.2,0.2))
for (i in 1:p){
  ts.plot(betas[,i],xlab="",main="",ylab="")
  mtext(side=3,line=0,paste("beta",i,sep=""))
}
ts.plot(sigma2s,xlab="",main="",ylab="")
mtext(side=3,line=0,"sigma2")

# Discarding burn-in draws
betas   = betas[seq(M0+1,niter,by=lag),]
sigma2s = sigma2s[seq(M0+1,niter,by=lag)]

# Autocorrelation of the retained draws
par(mfrow=c(3,3),mai=c(0.3,0.3,0.2,0.2))
for (i in 1:p){
  acf(betas[,i],main=paste("beta",i,sep=""),ylab="")
  mtext(side=3,line=0,paste("beta",i,sep=""))
}
acf(sigma2s,main="sigma2",ylab="")
mtext(side=3,line=0,"sigma2")

# Marginal posterior densities
par(mfrow=c(3,3),mai=c(0.3,0.3,0.2,0.2))
for (i in 1:p){
  hist(betas[,i],prob=TRUE,xlab=paste("beta",i,sep=""),main=namesX[i],breaks=25)
  curve(prior,from=-2,to=2,n=1000,col=4,lwd=2,add=TRUE)
  abline(v=bhat[i],col=2,lwd=2)
  abline(v=0,lwd=2,lty=2)
}
hist(sigma2s,prob=TRUE,xlab="",main="sigma2",breaks=30)
abline(v=s2,col=2,lwd=2)
curve(dgamma(x,a,b),from=0,to=0.5,n=1000,col=4,lwd=2,add=TRUE)