############################################################
#
# AR(1) plus noise model
#
# y(t) = x(t)       + e(t)        e(t) ~ N(0,V)
# x(t) = phi*x(t-1) + w(t)        w(t) ~ N(0,W)
#
# with E(e(t)w(l))=0 for all t,l
#
############################################################

sim.dlm = function(n,phi,V,W){
  sV  = sqrt(V)
  sW  = sqrt(W)
  x = rep(0,n)
  x[1] = rnorm(1,0,sW)
  for (t in 2:n)
    x[t] = rnorm(1,phi*x[t-1],sW)
  y = rnorm(n,x,sV)
  return(list(y=y,x=x))
}

ffbs = function(y,m0,c0,phi,V,W){
  n = length(y)	
  mf = rep(0,n)
  Cf = rep(0,n)
  m = m0
  C = C0
  for (t in 1:n){
    a = phi*m
    R = phi^2*C+W
    C = 1/(1/R+1/V)
    m = C*(a/R+y[t]/V)
    Cf[t] = C
    mf[t] = m
  }
  x  = rep(0,n)
  x[n] = rnorm(1,mf[n],sqrt(Cf[n]))
  for (t in (n-1):1){
    var = 1/(1/Cf[t]+phi^2/W)
    mean = var*(mf[t]/Cf[t]+phi*x[t+1]/W)
    x[t] = rnorm(1,mean,sqrt(var))
  }
  return(x)
}

pdf(file="AR1plusnoise-ffbs.pdf",width=12,height=10)

# Simulating some data
set.seed(12345)
n   = 200
V   = 1.00
W   = 0.25
phi = 0.95
sim = sim.dlm(n,phi,V,W)
y   = sim$y
x   = sim$x

par(mfrow=c(1,1))
ts.plot(y,type="b",ylab="",main="y(t) vs x(t)")
lines(x,col=2,type="b")

# MCMC via FFBS
m0   =  0.0
C0   = 100.0
phi0 =  0.95
D0   =  100.0
n0   = 6
V0   = 1
nu0  = 6
W0   = 0.25

# Initial value
phi1 = 1.0
V1   = 1
W1   = 1

M0 = 10000
M = 2000
L = 1
niter = M0+M*L
draws = matrix(0,niter,n+3)
for (iter in 1:niter){
  # sampling the states
  x1 = ffbs(y,m0,C0,phi1,V1,W1)
  # sampling phi
  var = 1/(1/D0+sum(x1[1:(n-1)]^2)/W1)
  mean = var*(phi0/D0+sum(x1[1:(n-1)]*x1[2:n])/W1)
  phi1 = rnorm(1,mean,sqrt(var)) 
  # sampling V
  a = n0+n
  b = n0*V0 + sum((y-x1)^2)
  V1 = 1/rgamma(1,a/2,b/2)
  # sampling W
  a = nu0+(n-1)
  b = nu0*W0 + sum((x1[2:n]-phi1*x1[1:(n-1)])^2)
  W1 = 1/rgamma(1,a/2,b/2)
  draws[iter,] = c(phi1,V1,W1,x1) 
}
ind = seq((M0+1),niter,by=L)

par(mfrow=c(3,3))
ts.plot(draws[,1],xlab="iterations",ylab="",main=expression(phi))
acf(draws[ind,1],main="")
hist(draws[ind,1],prob=TRUE,main="",xlab="")
abline(v=phi,col=2)
ts.plot(draws[,2],xlab="iterations",ylab="",main="V")
acf(draws[ind,2],main="")
hist(draws[ind,2],prob=TRUE,main="",xlab="")
abline(v=V,col=2)
ts.plot(draws[,3],xlab="iterations",ylab="",main="W")
acf(draws[ind,3],main="")
hist(draws[ind,3],prob=TRUE,main="",xlab="")
abline(v=W,col=2)

qx = t(apply(draws[ind,4:(n+3)],2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,1))
ts.plot(qx,col=c(3,2,3))
points(x,pch=16,cex=0.5)

dev.off()