##################################################################################
#
# Fitting a simple dynamic regression model
#
###################################################################################

############################################################
# 1st case
############################################################

# Simulating the data
# -------------------
set.seed(12345)
n     = 300
beta1 = 4
beta2 = 1
beta3 = -1
sig2  = 4
x     = rnorm(n)
y     = rep(0,n)
sig   = sqrt(sig2)
y[1:(n/3)]   = beta1*x[1:(n/3)]   + rnorm(n/3,0,sig)
y[(n/3+1):(2*n/3)] = beta2*x[(n/3+1):(2*n/3)] + rnorm(n/3,0,sig)
y[(2*n/3+1):n] = beta3*x[(2*n/3+1):n] + rnorm(n/3,0,sig)

par(mfrow=c(2,2))
plot(x,y,col=c(rep(1,n/3),rep(2,n/3),rep(3,n/3)),pch=16)
plot(x,y,pch=16)
ts.plot(y)
ts.plot(x)


# Kalman recursions: filtering and smoothing
# ------------------------------------------
tau2 = 0.1
m0   = 0
C0   = 1
m    = m0
C    = C0
mf   = rep(0,n)
Cf   = rep(0,n)
a    = rep(0,n)
R    = rep(0,n)

# Forward filtering
for (t in 1:n){
  a[t]  = m
  R[t]  = C+tau2
  ft    = x[t]*a[t]
  Qt    = x[t]^2*R[t]+sig2
  At    = R[t]*x[t]/Qt
  m     = a[t] + At*(y[t]-ft)
  C     = R[t] - At^2*Qt
  mf[t] = m
  Cf[t] = C
}
# Backward smoothing
mb = mf
Cb = Cf
for (t in (n-1):1){
  Cb[t] = Cf[t] - Cf[t]/R[t+1]*(R[t+1]-Cb[t+1])/R[t+1]*Cf[t]
  mb[t] = mf[t] + Cf[t]/R[t+1]*(mb[t+1]-a[t+1])
}

qf = cbind(mf-2*sqrt(Cf),mf,mf+2*sqrt(Cf))
qb = cbind(mb-2*sqrt(Cb),mb,mb+2*sqrt(Cb))

par(mfrow=c(1,1))
ts.plot(qf,ylim=range(qf,qb),lwd=2,lty=c(2,1,2),ylab="beta(t)")
lines(qb[,1],col=2,lty=2,lwd=2)
lines(qb[,2],col=2,lty=1,lwd=2)
lines(qb[,3],col=2,lty=2,lwd=2)
segments(0,beta1,n/3,beta1,col=4,lwd=2)
segments(n/3,beta2,2*n/3,beta2,col=4,lwd=2)
segments(2*n/3,beta3,n,beta3,col=4,lwd=2)
abline(v=n/3,col=4,lty=2)
abline(v=2*n/3,col=4,lty=2)
legend("topright",legend=c("Forward filtering","Backward smoothing"),lty=1:2)




############################################################
# 2nd case
############################################################


# Simulating the data
# -------------------
set.seed(54321)
n     = 300
sig2  = 4
tau2  = 0.1
x     = rnorm(n)
beta  = rep(0,n)
for (t in 2:n)
  beta[t] = beta[t-1]+rnorm(1,0,sqrt(tau2))
y = beta*x+rnorm(n,0,sqrt(sig2))

par(mfrow=c(1,2))
ts.plot(beta)
plot(x,y)

# Kalman recursions: filtering and smoothing
# ------------------------------------------
tau2 = 0.1
m0   = 0
C0   = 1
m    = m0
C    = C0
mf   = rep(0,n)
Cf   = rep(0,n)
a    = rep(0,n)
R    = rep(0,n)

# Forward filtering
for (t in 1:n){
  a[t]  = m
  R[t]  = C+tau2
  ft    = x[t]*a[t]
  Qt    = x[t]^2*R[t]+sig2
  At    = R[t]*x[t]/Qt
  m     = a[t] + At*(y[t]-ft)
  C     = R[t] - At^2*Qt
  mf[t] = m
  Cf[t] = C
}
# Backward smoothing
mb = mf
Cb = Cf
for (t in (n-1):1){
  Cb[t] = Cf[t] - Cf[t]/R[t+1]*(R[t+1]-Cb[t+1])/R[t+1]*Cf[t]
  mb[t] = mf[t] + Cf[t]/R[t+1]*(mb[t+1]-a[t+1])
}

qf = cbind(mf-2*sqrt(Cf),mf,mf+2*sqrt(Cf))
qb = cbind(mb-2*sqrt(Cb),mb,mb+2*sqrt(Cb))

par(mfrow=c(1,1))
ts.plot(qf,ylim=range(qf,qb),lwd=2,lty=c(2,1,2),ylab="beta(t)")
lines(qb[,1],col=2,lty=2,lwd=2)
lines(qb[,2],col=2,lty=1,lwd=2)
lines(qb[,3],col=2,lty=2,lwd=2)
lines(beta,col=4,lwd=3)
legend("topright",legend=c("True state","Forward filtering","Backward smoothing"),lty=1,col=c(4,1:2),lwd=2)