############################################################
#
# 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))
}

dlm = function(y,m0,c0,phi,V,W,moments){
  n = length(y)	
  mf = rep(0,n)
  Cf = rep(0,n)
  af = rep(0,n)
  Rf = rep(0,n)
  m = m0
  C = C0
  loglike = 0.0
  for (t in 1:n){
    a = phi*m
    R = phi^2*C+W
    loglike = loglike + dnorm(y[t],a,sqrt(R+V),log=TRUE)
    C = 1/(1/R+1/V)
    m = C*(a/R+y[t]/V)
    af[t] = a
    Rf[t] = R
    Cf[t] = C
    mf[t] = m
  }
  if (moments==1){
    mb = rep(0,n)
    Cb = rep(0,n)
    mb[n] = mf[n]
    Cb[n] = Cf[n]
    for (t in (n-1):1){
      B = phi*Cf[t]/Rf[t+1]
      Cb[t] = Cf[t] - B^2*(Rf[t+1]-Cb[t+1])
      mb[t] = mf[t] + B*(mb[t+1]-af[t+1])
    }
    return(list(loglike=loglike,mf=mf,Cf=Cf,mb=mb,Cb=Cb))
  }else{
    return(loglike)
  }
}


pdf(file="AR1plusnoise-Kalmanfilter-kalmansmoother.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

#y[15:16] = -10
ts.plot(y,type="b",ylab="",main="y(t) vs x(t)")
lines(x,col=2,type="b")

# Running both KF and KS for fixed (phi,V,W)
m0  =  0.0
C0  = 10.0
run = dlm(y,m0,c0,phi,V,W,moments=1)
mf = run$mf
mb = run$mb
Cf = run$Cf
Cb = run$Cb
sd.Cf = sqrt(Cf)
sd.Cb = sqrt(Cb)
limf = cbind(mf-2*sqrt(Cf),mf,mf+2*sqrt(Cf))
limb = cbind(mb-2*sqrt(Cb),mb,mb+2*sqrt(Cb))


par(mfrow=c(1,1))
plot(y,type="b",xlab="Time",ylab="")
lines(mf,col=2,type="b")
lines(mb,col=3,type="b")
legend("bottomright",legend=c(expression(y[t]),expression(m[t]^f),expression(m[t]^b)),col=1:3,lty=1)


par(mfrow=c(1,1))
plot(sd.Cf,xlab="time",ylab="Standard deviation",type="b",main="",ylim=range(sd.Cf,sd.Cb))
lines(sd.Cb,col=2,type="b")
legend("topright",legend=c(expression(C[t]^f),expression(C[t]^b)),col=1:2,lty=1)

par(mfrow=c(1,1))
plot(limf[,2],ylim=range(limf,limb),ylab="",type="l",xlab="Time")
lines(limf[,1])
lines(limf[,3])
lines(limb[,1],col=2)
lines(limb[,2],col=2)
lines(limb[,3],col=2)
legend("topright",legend=c("Kalman filter","Kalman smoother"),lty=1:2,col=1:2)
title("x(t)|y(1),...,y(t),theta vs x(t)|y(1),...,y(n),theta\ntheta=(phi,V,W)")


# Learning about phi (still fixing V and W)
N = 100
phis = seq(0.9,1.05,length=N)
loglike = rep(0,N)
for (i in 1:N)
  loglike[i] = dlm(y,m0,c0,phis[i],V,W,moments=0)
like = exp(loglike-max(loglike))

phi.hat = phis[like==max(like)]

plot(phis,like,type="l",ylab="Likelihood",xlab=expression(phi))
abline(v=phi,col=2)
abline(v=phi.hat,col=4)
legend("topleft",legend=c(paste("phi=",phi,sep=""),paste("phihat=",round(phi.hat,4),sep="")),col=c(2,4),lty=1)


# Comparing KF and KS for true phi and estimated phi
run1  = dlm(y,m0,c0,phi.hat,V,W,moments=1)
limf1 = cbind(run1$mf-2*sqrt(run1$Cf),run1$mf,run1$mf+2*sqrt(run1$Cf))
limb1 = cbind(run1$mb-2*sqrt(run1$Cb),run1$mb,run1$mb+2*sqrt(run1$Cb))

par(mfrow=c(1,1))
plot(limf[,2],ylim=range(limf,limf1),ylab="",type="l",xlab="Time")
lines(limf[,1])
lines(limf[,3])
lines(limf1[,1],col=2)
lines(limf1[,2],col=2)
lines(limf1[,3],col=2)
legend("topright",legend=c("KF - true phi","KF - MLE phi"),lty=1,col=1:2)

par(mfrow=c(1,1))
plot(limb[,2],ylim=range(limb,limb1),ylab="",type="l",xlab="Time")
lines(limb[,1])
lines(limb[,3])
lines(limb1[,1],col=2)
lines(limb1[,2],col=2)
lines(limb1[,3],col=2)
legend("topright",legend=c("KS - true phi","KS - MLE phi"),lty=1,col=1:2)


# Learning phi,V and W jointly
N = 50
phis = seq(0.9,1.05,length=N)
Vs   = seq(0.5,1.5,length=N)
Ws   = seq(0.1,0.5,length=N)
loglike = array(0,c(N,N,N))
for (i in 1:N)
  for (j in 1:N)
    for (k in 1:N)
      loglike[i,j,k] = dlm(y,m0,c0,phis[i],Vs[j],Ws[k],moments=0)
like = exp(loglike-max(loglike))
ind = 1:N
ii=ind[apply(like==max(like),1,sum)==1]
jj=ind[apply(like==max(like),2,sum)==1]
kk=ind[apply(like==max(like),3,sum)==1]
phi.hat = phis[ii]
V.hat = Vs[jj]
W.hat = Ws[kk]

contour(phis,Vs,like[,,kk],xlab="phi",ylab="V",drawlabels=FALSE)
points(phi,V,pch=16,col=2)
points(phi.hat,V.hat,pch=16,col=4)
legend("bottomleft",legend=c(paste("True - (",phi,",",V,")",sep=""),paste("MLE - (",round(phi.hat,3),",",round(V.hat,3),")",sep="")),col=c(2,4),pch=16)
title(paste("Likelihood of (phi,V)\n W=",round(W.hat,3),sep=""))


contour(phis,Ws,like[,jj,],xlab="phi",ylab="W",drawlabels=FALSE)
points(phi,W,pch=16,col=2)
points(phi.hat,W.hat,pch=16,col=4)
legend("bottomleft",legend=c(paste("True - (",phi,",",W,")",sep=""),paste("MLE - (",round(phi.hat,3),",",round(W.hat,3),")",sep="")),col=c(2,4),pch=16)
title(paste("Likelihood of (phi,W)\n V=",round(V.hat,3),sep=""))


contour(Vs,Ws,like[ii,,],xlab="V",ylab="W",drawlabels=FALSE)
points(V,W,pch=16,col=2)
points(V.hat,W.hat,pch=16,col=4)
legend("bottomleft",legend=c(paste("True - (",V,",",W,")",sep=""),paste("MLE - (",round(V.hat,3),",",round(W.hat,3),")",sep="")),col=c(2,4),pch=16)
title(paste("Likelihood of (V,W)\n phi=",round(phi.hat,3),sep=""))


# Comparing KF and KS for true (phi,V,W) versus estimated ones
run2  = dlm(y,m0,c0,phi.hat,V.hat,W.hat,moments=1)
limf2 = cbind(run2$mf-2*sqrt(run2$Cf),run2$mf,run2$mf+2*sqrt(run2$Cf))
limb2 = cbind(run2$mb-2*sqrt(run2$Cb),run2$mb,run2$mb+2*sqrt(run2$Cb))

par(mfrow=c(1,1))
plot(limf[,2],ylim=range(limf,limf2),ylab="",type="l",xlab="Time")
lines(limf[,1])
lines(limf[,3])
lines(limf2[,1],col=2)
lines(limf2[,2],col=2)
lines(limf2[,3],col=2)
legend("topright",legend=c("KF - true theta","KF - MLE theta"),lty=1,col=1:2)

par(mfrow=c(1,1))
plot(limb[,2],ylim=range(limb,limb2),ylab="",type="l",xlab="Time")
lines(limb[,1])
lines(limb[,3])
lines(limb2[,1],col=2)
lines(limb2[,2],col=2)
lines(limb2[,3],col=2)
legend("topright",legend=c("KS - true theta","KS - MLE theta"),lty=1,col=1:2)

dev.off()