###################################################################
#
#  Unit 10 - Example ii.
#
#  1st order dynamic linear model by the Forward filtering, 
#  backward sampling (FFBS) algorithm
#
###################################################################
#
# Author : Hedibert Freitas Lopes
#          Graduate School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@ChicagoGSB.edu
#
###################################################################
#  y(t)     ~ N(alpha(t)+F(t)*theta(t);V(t))
#  theta(t) ~ N(gamma+G*theta(t-1);W)                
ffbs = function(y,F,alpha,V,G,gamma,W,a1,R1,nd=1){
  n = length(y)
  if (length(F)==1){F = rep(F,n)}
  if (length(alpha)==1){alpha = rep(alpha,n)}
  if (length(V)==1){V = rep(V,n)}
  a = rep(0,n)
  R = rep(0,n)
  m = rep(0,n)
  C = rep(0,n)
  B = rep(0,n-1)
  H = rep(0,n-1)
  # time t=1
  a[1] = a1
  R[1] = R1
  f    = alpha[1]+F[1]*a[1]
  Q    = R[1]*F[1]**2+V[1]
  A    = R[1]*F[1]/Q
  m[1] = a[1]+A*(y[1]-f)
  C[1] = R[1]-Q*A**2
  # forward filtering
  for (t in 2:n){
    a[t] = gamma + G*m[t-1]
    R[t] = C[t-1]*G**2 + W
    f    = alpha[t]+F[t]*a[t]
    Q    = R[t]*F[t]**2+V[t]
    A    = R[t]*F[t]/Q
    m[t] = a[t]+A*(y[t]-f)
    C[t] = R[t]-Q*A**2
    B[t-1] = C[t-1]*G/R[t]
    H[t-1] = sqrt(C[t-1]-R[t]*B[t-1]**2)
  }
  # backward sampling
  theta = matrix(0,nd,n)
  theta[,n] = rnorm(nd,m[n],sqrt(C[n]))
  for (t in (n-1):1)
    theta[,t] = rnorm(nd,m[t]+B[t]*(theta[,t+1]-a[t+1]),sqrt(H[t]))
  if (nd==1){
    theta[1,]
  }
  else{
    theta
  }
}

quant05 = function(x){quantile(x,0.05)}
quant95 = function(x){quantile(x,0.95)}

######################################################
# Simulating a DLM                                   #
######################################################
set.seed(1576)
F = 1
G = 1
V = 0.25
W = 0.1
n = 100
theta = rep(0,n)
y     = rep(0,n)
theta[1] = 0
y[1] = rnorm(1,theta[1],sqrt(V))
for (t in 2:n){
  theta[t] = rnorm(1,theta[t-1],sqrt(W))
  y[t] = rnorm(1,theta[t],V)
}

par(mfrow=c(3,1))
ts.plot(y,main="y(t) vs true theta(t)",
        ylim=c(min(y,theta),max(y,theta)),type="b")
lines(theta,col=2)

thetas = ffbs(y,F,0.0,V,G,0.0,W,0,100,100)
mth = apply(thetas,2,median) 
q05 = apply(thetas,2,quant05) 
q95 = apply(thetas,2,quant95)
plot(theta,type="l",main="draws=100",ylim=c(min(q05),max(q95)))
lines(mth,col=4)
lines(q05,col=4)
lines(q95,col=4)

thetas = ffbs(y,F,0.0,V,G,0.0,W,0,100,10000)
mth = apply(thetas,2,median) 
q05 = apply(thetas,2,quant05) 
q95 = apply(thetas,2,quant95)
plot(theta,type="l",main="draws=10000",ylim=c(min(q05),max(q95)))
lines(mth,col=4)
lines(q05,col=4)
lines(q95,col=4)