####################################################################################
#
# Bayesian Gaussian AR(p) modeling of the Canadian Lynx time series 
#
# Our goal here is to compare MCMC schemes
#
# Parametes: beta0,beta1,...,betap,sigma2
#
# A) Gibbs sampler block moves: (beta0,...,betap) & sigma2 
#    2 blocks: 1st block with p+1 component, 2nd block with 1 component
# B) Gibbs sampler component-wise moves: beta0,beta1,...,betap,sigma2 
#    (p+2) single components
# C) RW-MH+Gibbs sampler component-wise moves: beta0,beta1,...,betap,sigma2 
#    (p+2) single components and RW-MH steps for beta0,...,betap.
#
# Lopes and Salazar (2004)
# Bayesian model uncertainty in smooth transition autoregressions
# Journal of Time Series Analysis, 27, 99-117.
# http://hedibert.org/wp-content/uploads/2013/12/lopes-salazar-2006a.pdf
#
# Hedibert F. Lopes
# www.hedibert.org
# February 3rd 2021
#
####################################################################################
library("astsa")

y = Lynx
y = y-mean(y)
y = y/sqrt(var(y))

par(mfrow=c(1,1))
ts.plot(y,main="Canadian Lynx",ylab="")

# Running the model with p=3
# --------------------------
data = y
p    = 3
nd   = length(data)
y    = data[(p+1):nd]
X    = cbind(1,data[p:(nd-1)])
for (i in 2:p)
  X = cbind(X,data[(p-i+1):(nd-i)])
n = nrow(X)

# prior hyperparameters
b0       = rep(0,p+1)
B0       = diag(25,p+1)
nu0      = 18
sig20    = 0.01666667

# posterior hyperparameters
nu0sig20 = nu0*sig20
iB0      = solve(B0)
iB0b0    = iB0%*%b0
B1       = solve(iB0+t(X)%*%X)
tcB1     = t(chol(B1))
b1       = B1%*%(iB0b0+t(X)%*%y)
nu1      = nu0 + n
sig21    = ((nu0sig20+sum((y-X%*%b1)*y)+t(b0-b1)%*%iB0b0)/nu1)[1]
se       = sqrt(nu1/(nu1-2)*sig21*diag(B1))

tab = round(cbind(b1,se,b1/se,2*(1-pt(abs(b1/se),df=nu1))),5)
colnames(tab)=c("Mean","StDev","t","tail")
rownames(tab)=c("Intercept","beta1","beta2","beta3")
tab

# Gibbs sampler: components (beta and sigma2)
# -------------------------------------------
set.seed(12345)
M      = 2000
burnin = 10000
skip   = 10
niter  = burnin+skip*M
draws  = matrix(0,niter,p+2)
beta.d = rep(0,p+1)
runtime = system.time({
for (iter in 1:niter){
  sig2.d = 1/rgamma(1,nu1/2,(nu0sig20+sum((y-X%*%beta.d)^2))/2)
  beta.d = b1+sqrt(sig2.d)*tcB1%*%rnorm(p+1)
  draws[iter,] = c(beta.d,sig2.d)
}})
draws  = draws[seq(burnin+1,niter,by=skip),]

# Gibbs sampler: components (beta0,beta1,...,betap,sigma2)
# --------------------------------------------------------
ind = 1:(p+1)
set.seed(54321)
niter  = burnin+skip*M
draws1 = matrix(0,niter,p+2)
beta.d = rep(0,p+1)
runtime1 = system.time({
for (iter in 1:niter){
  sig2.d = 1/rgamma(1,nu1/2,(nu0sig20+sum((y-X%*%beta.d)^2))/2)
  for (i in 1:(p+1)){
  	XX = X[,ind!=i]
  	yy = y-XX%*%beta.d[ind!=i]
  	xx = X[,ind==i]
  	var = 1/(sum(xx^2)+1/B0[i,i])
  	mean = var*(sum(xx*yy)+b0[i]/B0[i,i])
  	beta.d[i] = rnorm(1,mean,sqrt(sig2.d*var))
  }
  draws1[iter,] = c(beta.d,sig2.d)
}})
draws1 = draws1[seq(burnin+1,niter,by=skip),]

# MCMC: cycling beta via RW-MH and Gibbs for sigma2
# -------------------------------------------------
ind = 1:(p+1)
set.seed(54321)
niter  = burnin+skip*M
draws2 = matrix(0,niter,p+2)
beta.d = rep(0,p+1)
runtime2 = system.time({
for (iter in 1:niter){
  sig2.d = 1/rgamma(1,nu1/2,(nu0sig20+sum((y-X%*%beta.d)^2))/2)
  for (i in 1:(p+1)){
  	XX = X[,ind!=i]
  	yy = y-XX%*%beta.d[ind!=i]
  	xx = X[,ind==i]
  	beta1 = rnorm(1,beta.d[i],0.01)
    num = sum(dnorm(yy,xx*beta1,sqrt(sig2.d),log=TRUE))
    den = sum(dnorm(yy,xx*beta.d[i],sqrt(sig2.d),log=TRUE))
    accept = min(0,num-den)
    if (log(runif(1))<accept){
      beta.d[i] = beta1
    }  	
  }
  draws2[iter,] = c(beta.d,sig2.d)
}})
draws2 = draws2[seq(burnin+1,niter,by=skip),]

runtime
runtime1
runtime2


names = c("beta0","beta1","beta2","beta3","sigma2")
par(mfrow=c(3,p+2),mai=c(0.35,0.35,0.2,0.1))
for (i in 1:(p+2)){
  range = range(draws[,i],draws1[,i],draws2[,i])
  ts.plot(draws[,i],main=names[i],ylim=range)
  lines(draws1[,i],col=2)  
  lines(draws2[,i],col=3)
}
for (i in 1:(p+2)){
  acf1 = acf(draws[,i],plot=FALSE,lag.max=200)
  acf2 = acf(draws1[,i],plot=FALSE,lag.max=200)
  acf3 = acf(draws2[,i],plot=FALSE,lag.max=200)
  range = range(acf1$acf,acf2$acf,acf3$acf)
  plot(acf1$lag,acf1$acf,type="l",ylim=range)
  lines(acf2$lag,acf2$acf,type="l",col=2)
  lines(acf3$lag,acf3$acf,type="l",col=3)
}
legend("topright",legend=c("Gibbs block","Gibbs single","RW-MH+Gibbs single"),col=1:3,lty=1,bty="n",lwd=2)


for (i in 1:(p+1)){
  den   = density(draws[,i])
  den1  = density(draws1[,i])
  den2  = density(draws2[,i])
  xxx   = seq(min(draws[,i]),max(draws[,i]),length=1000)
  den.e = dt((xxx-b1[i])/sqrt(sig21*B1[i,i]),df=nu1)/sqrt(sig21*B1[i,i])
  range  = range(den$y,den1$y,den.e)
  plot(den$x,den$y,xlab="",main=paste("beta ",i-1,sep=""),ylim=range,type="l")
  lines(den1$x,den1$y,col=2)
  lines(den2$x,den1$y,col=3)
  lines(xxx,den.e,col=4)
}
sig2.d  = draws[,p+2]
sig2.d1 = draws1[,p+2]
sig2.d2 = draws2[,p+2]
plot(density(sig2.d),xlab="",ylab="Posterior density",main="sigma2")
lines(density(sig2.d1),col=2)
lines(density(sig2.d2),col=3)
xxx = seq(min(sig2.d),max(sig2.d),length=1000)
lines(xxx,dgamma(1/xxx,nu1/2,nu1*sig21/2)/xxx^2,col=4)
legend("topright",legend="Exact posterior",col=4,lty=1,bty="n",lwd=2)


# Effective sample size (measuring the correlation of the chains)
ess = matrix(0,p+2,3)
for (i in 1:(p+2)){
  ess[i,1] = 1/(1+2*sum(acf(draws[,i],lag.max=200,plot=FALSE)$acf))
  ess[i,2] = 1/(1+2*sum(acf(draws1[,i],lag.max=200,plot=FALSE)$acf))
  ess[i,3] = 1/(1+2*sum(acf(draws2[,i],lag.max=200,plot=FALSE)$acf))
}

tab = round(100*ess,1)
rownames(tab) = names
colnames(tab) = c("Gibbs block","Gibbs single","RW-MH+Gibbs single")
M*tab/100

########################################################
#            M=2000 - burnin=10000 - skip=10
########################################################
#        Gibbs block   Gibbs single   RW-MH+Gibbs single
########################################################
# beta0          418            784                  352
# beta1          626            524                  284
# beta2          576            652                  134
# beta3          458            716                  138
# sigma2         858            904                  984
########################################################

########################################################
#      M      = 2000       M = 2000      M      =   2000
#      burnin = 1000  burnin = 1000      burnin = 100000
#      skip   =    1    skip =   10        skip =     50
########################################################
#        Gibbs block   Gibbs single   RW-MH+Gibbs single
########################################################
# beta0          782           738                   350
# beta1          826           414                   228
# beta2          630           486                   178
# beta3          528           502                   162
# sigma2         592           940                   800
########################################################