# kalman recursion
kalmanfilter = function(y,m0,c0,sig2,phi,tau2){
  n = length(y)
  m = m0
  C = C0
  mf = rep(m,n+1)
  Cf = rep(C,n+1)
  loglike = 0.0
  for (t in 1:n){
    a = phi*m
    R = phi^2*C+tau2
    loglike = loglike + dnorm(y[t],a,sqrt(R+sig2),log=TRUE)
    C = 1/(1/R+1/sig2)
    m = C*(a/R+y[t]/sig2)
    mf[t+1] = m
    Cf[t+1] = C
  }
  return(list(mf=mf,Cf=Cf,loglike=loglike))
}

# Simulatig the dataset
set.seed(1235)
n    = 100
phi  = 0.9
tau  = 0.5
sig  = 1.0
tau2 = tau^2
sig2 = sig^2
x    = rep(0,2*n)
for (t in 2:(2*n))
  x[t] = phi*x[t-1]+rnorm(1,0,tau)
x.true = x[(n+1):(2*n)]
y  = rnorm(n,x.true,sig)

par(mfrow=c(1,1))
plot(y,pch=16,xlab="time",ylab="")
lines(x.true,col=2)


# x0~N(m0,C0)
m0 = 0
C0 = 1
 
# Kalman momments
kf = kalmanfilter(y,m0,C0,sig2,phi,tau2)
mf = kf$mf
Cf = kf$Cf
qkf = cbind(mf+qnorm(0.05)*sqrt(Cf),mf,mf+qnorm(0.95)*sqrt(Cf))

par(mfrow=c(1,1))
plot(y,ylim=range(y,qkf),xlab="time",ylab="states",cex=0.75,pch=16,
     xlim=c(0,n))
lines(0:n,qkf[,2],col=2)
lines(0:n,qkf[,1],col=2,lty=2)
lines(0:n,qkf[,3],col=2,lty=2)
title("Kalman filter")

# Using the Kalman filter to compute the
# integrated likelihood for the pair (phi,tau)
N = 100
phis  = seq(0.0,1.1,length=N)
tau2s = seq(0.01,1.1,length=N)^2
loglike = matrix(0,N,N)
for (i in 1:N)
  for (j in 1:N)
    loglike[i,j] = kalmanfilter(y,m0,C0,sig2,phis[i],tau2s[j])$loglike

contour(phis,sqrt(tau2s),exp(loglike),xlab=expression(phi),
ylab=expression(tau),drawlabels=FALSE)
points(phi,tau,col=2,pch=16)
title("Likelihood function")

ind=1:N
ii=ind[apply(loglike==max(loglike),1,sum)==1]
jj=ind[apply(loglike==max(loglike),2,sum)==1]
loglike[ii,jj]
cbind(c(phi,tau),round(c(phis[ii],sqrt(tau2s[jj])),2))
points(phis[ii],sqrt(tau2s[jj]),pch=16)