############################################################################################
#
# Autoregressive model of order one, AR(1), with normal errors
#
#  Model: y(t) = rho*y(t-1) + e(t)   t=1,...,n
#
#  Prior:
#         y(0) ~ N(m0,C0)
#         rho  ~ N(r0,V0)
#         sig2 ~ IG(n0/2,n0*s02/2)
# 
############################################################################################
#
#  HEDIBERT FREITAS LOPES
#  Associate Professor of Econometrics and Statistics
#  The University of Chicago Booth School of Business
#  5807 South Woodlawn Avenue
#  Chicago, Illinois, 60637
#  Email : hlopes@ChicagoGSB.edu
#  URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################
q025 = function(x){quantile(x,0.025)}
q975 = function(x){quantile(x,0.975)}

# Simulating data
set.seed(1255)
n         = 100
y0.true   = 0.0
rho.true  = 0.95
sig2.true = 1.0
sig.true  = sqrt(sig2.true)
y         = rep(0,n)
y[1]      = rnorm(1,rho.true*y0.true,sig.true)
for (t in 2:n)
  y[t] = rnorm(1,rho.true*y[t-1],sig.true)


pdf(file="ar1-data.pdf",width=10,height=5)
par(mfrow=c(1,2))
plot(y,xlab="Time",ylab="",main=expression(y[t]),type="b")
plot(y[1:(n-1)],y[2:n],xlab=expression(y[t-1]),ylab=expression(y[t]),ylim=range(y),xlim=range(y))
abline(0,1,col=2,lwd=3)
dev.off()


# Prior hyperparameters
m0    =  0.0
C0    = 10.0
r0    =  0.9
V0    = 10.0
n0    =  5.0
s02   =  1.0
iC0   = 1/C0
iC0m0 = iC0*m0
iV0   = 1/V0
iV0r0 = iV0*r0
n0s02 = n0*s02
n1    = n0+n

# MCMC setup
M0    = 1000
M     = 1000
L     = 1
niter = M0+M*L
rho   = rho.true
sig2  = sig2.true
y0    = y0.true
pars  = NULL
for (i in 1:niter){
  x    = c(y0,y[1:(n-1)])
  v    = 1/(iV0+sum(x^2)/sig2)
  m    = v*(iV0r0+sum(x*y)/sig2)
  rho  = rnorm(1,m,sqrt(v))
  sig2 = rgamma(1,n1/2,(n0s02+sum((y-rho*x)^2))/2)
  v    = 1/(iC0+rho^2/sig2)
  m    = v*(iC0m0+rho*y[1]/sig2)
  y0   = rnorm(1,m,sqrt(v))
  pars = rbind(pars,c(rho,sig2))
}
pars = pars[seq(M0+1,niter,by=L),]

true = c(rho.true,sig2.true)
names = c("rho","sigma2")
pdf(file="ar1-mcmc.pdf",width=10,height=7)
par(mfrow=c(2,3))
for (i in 1:2){
  ts.plot(pars[,i],xlab="iteration",ylab="",main=names[i])
  abline(h=true[i],col=2,lwd=3)
  acf(pars[,i],main="")
  hist(pars[,i],xlab="",ylab="",main="")
  abline(v=true[i],col=2,lwd=3)
}
dev.off()

# k-steps ahead forecasting
rho = pars[,1]
sig = sqrt(pars[,2])
k   = 20
yf  = matrix(0,M,k)
yf[,1] = rnorm(M,rho*y[n],sig)
for (i in 2:k)
  yf[,i] = rnorm(M,rho*yf[,i-1],sig)
  
myf = apply(yf,2,median)
lyf = apply(yf,2,q025)
uyf = apply(yf,2,q975)

set.seed(12325)
ind = sample(1:M,size=100,replace=FALSE,prob=rep(1/M,M))
pdf(file="ar1-forecast.pdf",width=10,height=5)
par(mfrow=c(1,2))
ts.plot(t(yf[ind,]),col=1:20,ylim=range(yf))
title("100 forecasting paths")
plot(y,ylim=range(yf),xlim=c(1,n+k),xlab="Time",ylab="",main="",pch=16,cex=0.5)
lines((n+1):(n+k),lyf,col=1,lwd=3)
lines((n+1):(n+k),myf,col=1,lwd=3)
lines((n+1):(n+k),uyf,col=1,lwd=3)
title("95% forecasting interval")
dev.off()