###################################################################
#
#  Unit 10 - Example i.
#
#  1st order and 2nd order dynamic linear models
#
###################################################################
#
# Author : Hedibert Freitas Lopes
#          Graduate School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@ChicagoGSB.edu
#
###################################################################
set.seed(123456)

# Simulating 1st order DLM
# ------------------------
W = 0.5
V = 1.0
n = 100
x0 = 0.00
w  = rnorm(n,0,sqrt(W))
v  = rnorm(n,0,sqrt(V))
x  = y = rep(0,n)
x[1] = x0   + w[1]
y[1] = x[1] + v[1]
for (t in 2:n){
  x[t] = x[t-1] + w[t]
  y[t] = x[t] + v[t]
}
y1 = y
x1 = x

# Simulating 2nd order DLM
# ------------------------
# 2nd order 
W1 = 0.001
W2 = 0.001
V1 = 0.1
n = 100
x0 = c(0.00,0.00)
w  = cbind(rnorm(n,0,sqrt(W1)),rnorm(n,0,sqrt(W2)))
v  = rnorm(n,0,sqrt(V1))
y = rep(0,n)
x = matrix(0,n,2)
x[1,2] = x0[2]          + w[1,2]
x[1,1] = x0[1] + x[1,2] + w[1,1]
y[1]   = x[1,1]         + v[1]
for (t in 2:n){
  x[t,2] = x[t-1,2]          + w[t,2]
  x[t,1] = x[t-1,1] + x[t,2] + w[t,1]
  y[t]   = x[t,1]            + v[t]
}
y2 = y
x2 = x

# Plotting 1st order and 2nd order DLMs
# -------------------------------------
par(mfrow=c(2,1))
plot(y1,type="l",xlab="time",ylab="")
lines(x1,col=2)
title(paste("V=",V," - W=",W,sep=""))
plot(y2,type="l",xlab="time",ylab="")
lines(x2[,1],col=2)
lines(x2[,2],col=4)
title(paste("V=",V1," - W1=",W1," - W2=",W2,sep=""))

# Online fitting 1st order DLM
# -----------------------------
  y  = y1
  a1 = 0
  R1 = 100
  m  = rep(0,n)
  C  = rep(0,n)
  B  = rep(0,n-1)
  H  = rep(0,n-1)
  # time t=1
  a = a1
  R = R1
  f = a
  Q = R+V
  A = R/Q
  m[1] = a+A*(y[1]-f)
  C[1] = R-Q*A^2
  # forward filtering
  for (t in 2:n){
    a = m[t-1]
    R = C[t-1] + W
    f = a
    Q = R+V
    A = R/Q
    m[t] = a+A*(y[t]-f)
    C[t] = R-Q*A^2
}
m1 = m
C1 = C

# Online fitting 2nd order DLM
# ----------------------------
y  = y2
F  = matrix(c(1,0),2,1)
G  = matrix(c(1,0,1,1),2,2)
Ws = diag(c(W1,W2))
a1 = matrix(0,2,1)
R1 = diag(100,2)
a  = matrix(0,2,1)
R  = matrix(0,2,2)
m  = matrix(0,n,2)
C  = array(0,c(n,2,2))
# time t=1
a = a1
R = R1
f = t(F)%*%a
Q = t(F)%*%R%*%F+V1
A = R%*%F/Q[1,1]
m[1,]  = a+A%*%(y[1]-f)
C[1,,] = R-A%*%t(A)*Q[1,1]
# forward filtering
for (t in 2:n){
  a = G%*%m[t-1,]
  R = G%*%C[t-1,,]%*%t(G) + Ws
  f = t(F)%*%a
  Q = t(F)%*%R%*%F+V1
  A = R%*%F/Q[1,1]
  m[t,]  = a+A%*%(y[t]-f)
  C[t,,] = R-A%*%t(A)*Q[1,1]
}
m2 = m
C2 = C

# Plotting fitted 1st order and 2nd order DLMs
# --------------------------------------------
par(mfrow=c(3,1))
L  = m1-2*C1
U  = m1+2*C1
ts.plot(x1,xlab="time",ylab="",ylim=c(min(L),max(U)))
lines(x1,col=1,lwd=1.5)
lines(m1,col=4,lwd=1.5)
lines(L,col=4)
lines(U,col=4)
title("1st order DLM")
names = c("intercept","slope")
for (i in 1:2){
  L = m2[,i]-2*sqrt(C2[t,i,i])
  U = m2[,i]+2*sqrt(C2[t,i,i])
  ts.plot(x2[,i],xlab="time",ylab="",ylim=c(min(L),max(U)))
  lines(x2[,i],col=1,lwd=1.5)
  lines(m2[,i],col=4,lwd=1.5)
  lines(L,col=4)
  lines(U,col=4)
  title(paste("2nd order DLM - ",names[i],sep=""))
}

# Smoothing the 1st order DLM
# ---------------------------
  y  = y1
  a1 = 0
  R1 = 100
  a  = rep(0,n)
  R  = rep(0,n)
  m  = rep(0,n)
  C  = rep(0,n)
  mm = rep(0,n-1)
  CC = rep(0,n-1)
  # time t=1
  a[1] = a1
  R[1] = R1
  f = a[1]
  Q = R[1]+V
  A = R[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] = m[t-1]
    R[t] = C[t-1] + W
    f = a[t]
    Q = R[t]+V
    A = R[t]/Q
    m[t] = a[t]+A*(y[t]-f)
    C[t] = R[t]-Q*A^2
  }
  mm[n] = m[n]
  CC[n] = C[n]
  for (t in (n-1):1){
    mm[t] = m[t] + C[t]/R[t+1]*(mm[t+1]-a[t+1])
    CC[t] = C[t] - (C[t]^2)/(R[t+1]^2)*(R[t+1]-CC[t+1])
  }

par(mfrow=c(2,1))
L1  = m-2*C
U1  = m+2*C
L2  = mm-2*CC
U2  = mm+2*CC
ts.plot(x1,xlab="time",ylab="",ylim=c(min(L1,L2),max(U1,U2)))
lines(x1,col=1,lwd=1.5)
lines(m,col=4,lwd=1.5)
lines(L1,col=4)
lines(U1,col=4)
title("Online distributions")
ts.plot(x1,xlab="time",ylab="",ylim=c(min(L1,L2),max(U1,U2)))
lines(x1,col=1,lwd=1.5)
lines(mm,col=4,lwd=1.5)
lines(L2,col=4)
lines(U2,col=4)
title("Smoothed distributions")