rm(list=ls())
############################
# Simulating the data
############################
set.seed(469479)
n    = 100
sig2 = 1
tau2 = 0.25
x    = rep(0,n)
y    = rep(0,n)
y[1] = x[1]+rnorm(1,0,sqrt(sig2))
for (t in 2:n){
  x[t] = x[t-1]+rnorm(1,0,sqrt(tau2))
  y[t] = x[t]+rnorm(1,0,sqrt(sig2))
}
x.true = x

pdf(file="toy-dlm.pdf",width=20,height=15)
par(mfrow=c(1,1))
plot(y,pch=16)
title("y(t) versus x(t) - 1st order DLM")
lines(x)

############################
# Kalman filter
############################
m0 = 0
C0 = 10
m  = m0
C  = C0
stat = matrix(0,n,2)
B  = rep(0,n)
mf = rep(0,n)
Cf = rep(0,n)
for (t in 1:n){
  R = C + tau2
  Q = R + sig2
  A = R/Q
  m = (1-A)*m+A*y[t]
  C = A*sig2
  mf[t] = m
  Cf[t] = C
  B[t] = C/(C+tau2)
}
qf = cbind(mf-2*sqrt(Cf),mf,mf+2*sqrt(Cf))

par(mfrow=c(1,1))
plot(x,ylim=range(x,qf),xlab="Time",ylab="",pch=16,type="b")
lines(qf[,1],col=2)
lines(qf[,2],col=2)
lines(qf[,3],col=2)
title("True x(t) versus Kalman filter")

############################
# Kalman Smoother
############################
mb = rep(0,n)
Cb = rep(0,n)
mb[n] = mf[n]
Cb[n] = Cf[n]
for (t in (n-1):1){
  mb[t] = mf[t] + B[t]*(mb[t+1]-mf[t])
  Cb[t] = Cf[t]-B[t]^2*(Cf[t]+tau2-Cb[t+1])
}
qb = cbind(mb-2*sqrt(Cb),mb,mb+2*sqrt(Cb))

par(mfrow=c(1,1))
plot(qf[,1],ylim=range(x,qf,qb),col=2,type="l",xlab="Time",ylab="")
title("Kalman filter (red) versus Kalman smoother (green)")
lines(qf[,2],col=2)
lines(qf[,3],col=2)
lines(qb[,1],col=3)
lines(qb[,2],col=3)
lines(qb[,3],col=3)

############################
# Kalman Sampler
############################
M      = 10000
ab     = matrix(0,M,n)
Rb     = matrix(0,M,n)
ab[,n] = mf[n]
Rb[,n] = Cf[n]
xdraw  = matrix(0,M,n)
xdraw[,n] = rnorm(M,ab[,n],sqrt(Rb[,n]))
for (t in (n-1):1){
  ab[,t] = (1-B[t])*mf[t]+B[t]*xdraw[,t+1]
  Rb[,t] = Cf[t] - B[t]^2*(Cf[t]+tau2)
  xdraw[,t] = rnorm(M,ab[,t],sqrt(Rb[,t]))
}
qd = t(apply(xdraw,2,quantile,c(0.025,0.5,0.975)))
ts.plot(qd,ylim=range(qd,qb),col=2)
lines(qb[,1],col=3)
lines(qb[,2],col=3)
lines(qb[,3],col=3)
title("Kalman smoother (green) versus posterior draws (red)")


############################################# 
# GIBBS SAMPLER: sampling (x1,...,xn) jointly
############################################# 
m0 = 0
C0 = 1000
a  = 0.001
b  = 0.001
c  = 0.001
d  = 0.001
sig2draw = 1
tau2draw = 0.25
burn  = 1000
niter = burn+1000
draws = matrix(0,niter,n+2)
for (iter in 1:niter){
# Sample x's
m = m0
C = C0
mf = rep(0,n)
Cf = rep(0,n)
B = rep(0,n)
for (t in 1:n){
  R = C + tau2draw
  Q = R + sig2draw
  A = R/Q
  m = (1-A)*m+A*y[t]
  C = A*sig2draw
  mf[t] = m
  Cf[t] = C
  B[t] = C/(C+tau2draw)
}
ab    = rep(0,n)
Rb    = rep(0,n)
ab[n] = mf[n]
Rb[n] = Cf[n]
xdraw = rep(0,n)
xdraw[n] = rnorm(1,ab[n],sqrt(Rb[n]))
for (t in (n-1):1){
  ab[t] = (1-B[t])*mf[t]+B[t]*xdraw[t+1]
  Rb[t] = Cf[t] - B[t]^2*(Cf[t]+tau2draw)
  xdraw[t] = rnorm(1,ab[t],sqrt(Rb[t]))
}

# Sample sig2 and tau2
sig2draw = 1/rgamma(1,a+n/2,b+sum((y-xdraw)^2)/2)
tau2draw = 1/rgamma(1,c+(n-1)/2,d+sum(diff(xdraw)^2)/2)

# Storing draws
draws[iter,] = c(xdraw,sig2draw,tau2draw)
}
draws= draws[(burn+1):niter,]
quant = t(apply(draws,2,quantile,c(0.025,0.5,0.975)))

par(mfrow=c(1,2))
plot(draws[,(n+1):(n+2)],xlab=expression(sigma^2),ylab=expression(tau^2),pch=16)
points(sig2,tau2,pch=16,cex=3,col=2)
title("Posterior draws of (sig2,tau2)\nBased on FFBS")
ts.plot(quant[1:n,],xlab="Time",ylab="",ylim=range(quant[1:n,],qb,y),col=4)
points(y,pch=16)
title("Posterior quantiles of (x1,...,xn)")

################################################
# GIBBS SAMPLER: sampling x1,...,xn individually
################################################
m0       = 0
C0       = 1000
a        = 0.001
b        = 0.001
c        = 0.001
d        = 0.001
sig2draw = 1
tau2draw = 0.25
burn     = 1000
niter    = burn+1000
draws1   = matrix(0,niter,n+2)
xdraw = qb[,2]
for (iter in 1:niter){
  # Sample x's
  R    = C0+tau2draw
  var  = 1/(1/sig2draw+1/tau2draw+1/R)
  mean = var*(y[1]/sig2draw+1/tau2draw+m0/R)
  xdraw[1] = rnorm(1,mean,sqrt(var))
  for (t in 2:(n-1)){
	  var      = 1/(1/sig2draw+2/tau2draw)
	  mean     = var*(y[t]/sig2draw+(xdraw[t-1]+xdraw[t+1])/tau2draw)
	  xdraw[t] = rnorm(1,mean,sqrt(var))
  }
  var      = 1/(1/sig2draw+1/tau2draw)
  mean     = var*(y[n]/sig2draw+xdraw[n-1]/tau2draw)
  xdraw[n] = rnorm(1,mean,sqrt(var))
  # Sample sig2 and tau2
  sig2draw = 1/rgamma(1,a+n/2,b+sum((y-xdraw)^2)/2)
  tau2draw = 1/rgamma(1,c+(n-1)/2,d+sum(diff(xdraw)^2)/2)
  # Storing draws
  draws1[iter,] = c(xdraw,sig2draw,tau2draw)
}
draws1 = draws1[(burn+1):niter,]
quant1 = t(apply(draws1,2,quantile,c(0.025,0.5,0.975)))

par(mfrow=c(2,2))
acfg = acf(draws[,1],plot=FALSE)
plot(acfg$lag,acfg$acf,type="l",ylim=c(-0.2,1),xlab="LAG",ylab="ACF",main="Joint move (FFBS)")
for (t in 2:n){
  acfg = acf(draws[,t],plot=FALSE)
  lines(acfg$lag,acfg$acf)
}
acfg = acf(draws1[,1],plot=FALSE)
plot(acfg$lag,acfg$acf,type="l",ylim=c(-0.2,1),xlab="LAG",ylab="ACF",main="Individual move",col=2)
for (t in 2:n){
  acfg = acf(draws1[,t],plot=FALSE)
  lines(acfg$lag,acfg$acf,col=2)
}

ts.plot(quant[1:n,],xlab="Time",ylab="",ylim=range(quant[1:n,],quant1[1:n,]))
title("Posterior quantiles of (x1,...,xn)")
lines(quant1[1:n,1],col=2)
lines(quant1[1:n,2],col=2)
lines(quant1[1:n,3],col=2)

plot(draws[,(n+1):(n+2)],xlab=expression(sigma^2),ylab=expression(tau^2),
     xlim=range(draws[,(n+1)],draws1[,(n+1)]),ylim=range(draws[,(n+2)],draws1[,(n+2)]))
points(draws1[,(n+1):(n+2)],col=2)
points(sig2,tau2,cex=3,pch=16,col=3)
title("Posterior draws of (sig2,tau2)")
dev.off()