#######################################################
#
# AR(1) plus noise
# 
#  y(t) = x(t)       + v(t)          v(t) iid N(0,sig2)
#  x(t) = phi*x(t-1) + w(t)          w(t) iid N(0,tau2) 
#
# with x(0) ~ N(m0,C0).
#
#######################################################

#######################################################
# kalman filtering and backward sampling
#######################################################
kalmanfilter = function(y,m0,c0,sig2,phi,tau2,M){
  n = length(y)
  m = m0
  C = C0
  mf = rep(m,n)
  Cf = rep(C,n)
  loglike = 0.0
  x  = matrix(0,n,M)
  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] = m
    Cf[t] = C
  }
  x[n,] = rnorm(M,mf[n],sqrt(Cf[n])) 
  for (t in (n-1):1){
    var = 1/(phi^2/tau2+1/Cf[t])
    mean=var*(phi*x[t+1,]/tau2+mf[t]/Cf[t])
    x[t,] = rnorm(M,mean,sqrt(var))
  }
  return(list(mf=mf,Cf=Cf,loglike=loglike,x=x))
}


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

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


# x0~N(m0,C0)
m0 = 0
C0 = 1
 
#######################################################
# Kalman momments (conditional on phi,tau2,sig2)
#######################################################
kf = kalmanfilter(y,m0,C0,sig2,phi,tau2,10000)
mf = kf$mf
Cf = kf$Cf
qkf = cbind(mf+qnorm(0.05)*sqrt(Cf),mf,mf+qnorm(0.95)*sqrt(Cf))
qks = t(apply(kf$x,1,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(1,1))
plot(qkf[,2],ylim=range(qkf,qks),xlab="time",ylab="states",col=2,type="l")
lines(qkf[,1],col=2,lty=2)
lines(qkf[,3],col=2,lty=2)
lines(qks[,2],col=4)
lines(qks[,1],col=4,lty=2)
lines(qks[,3],col=4,lty=2)
title("Kalman filtering and smoothing\n Conditional on phi,tau2,sig2")
legend("topright",legend=c("Kalman filter","Kalman smoother"),col=c(2,4),lty=1)

#######################################################
# Bayesian inference via MCMC (FFBS)
#######################################################
# Prior for (phi,tau2,sig2)
# phi  ~ N(f0,F0)
# tau2 ~ IG(nu0/2,nu0*tau20/2)
# sig2 ~ IG(eta0/2,eta0*sig20/2)
f0=0
F0=100
nu0=1
tau20=1
eta0=
sig20=1

# Initial valuess
set.seed(2325)
tau2 = 1.0
sig2 = 1.0
x    = kalmanfilter(y,m0,C0,sig2,phi,tau2,1)$x

# MCMC setup
M0    = 1000
M     = 2000
L     = 10
niter = M0+M*L
draws = matrix(0,niter,n+3)
for (iter in 1:niter){
  # Sampling phi
  var = 1/(1/F0+sum(x[1:(n-1),]^2)/tau2)
  mean = var*(f0/F0+sum(x[1:(n-1)]*x[2:n])/tau2)
  phi = rnorm(1,mean,sqrt(var))

  # Sampling tau2
  par1 = nu0+(n-1)
  par2 = nu0*tau20+sum((x[2:n]-phi*x[1:(n-1)])^2)
  tau2 = 1/rgamma(1,par1/2,par2/2)

  # Sampling sig2
  par1 = eta0+n
  par2 = eta0*sig20+sum((y-x)^2)
  sig2 = 1/rgamma(1,par1/2,par2/2)

  # Sampling the states via FFBS
  x = kalmanfilter(y,m0,C0,sig2,phi,tau2,1)$x
  draws[iter,]=c(x,phi,sqrt(tau2),sqrt(sig2))
}

ind = seq(M0+1,niter,by=L)
par(mfrow=c(3,3))
ts.plot(draws[,n+1],xlab="iteration",ylab="",main=expression(phi))
abline(h=phi.t,col=2,lwd=3)
ts.plot(draws[,n+2],xlab="iteration",ylab="",main=expression(tau))
abline(h=tau.t,col=2,lwd=3)
ts.plot(draws[,n+3],xlab="iteration",ylab="",main=expression(sigma))
abline(h=sig.t,col=2,lwd=3)
for (i in 1:3)
  acf(draws[ind,n+i],main="")
for (i in 1:3){
  hist(draws[ind,n+i],main="",xlab="",prob=TRUE)
  abline(v=true[i],col=2,lwd=3)
}

par(mfrow=c(1,3))
plot(draws[ind,n+1],draws[ind,n+2],xlab=expression(phi),ylab=expression(tau))
points(phi.t,tau.t,col=2,pch=16,cex=2)
plot(draws[ind,n+1],draws[ind,n+3],xlab=expression(phi),ylab=expression(sigma))
points(phi.t,sig.t,col=2,pch=16,cex=2)
plot(draws[ind,n+2],draws[ind,n+3],xlab=expression(taus),ylab=expression(sigma))
points(tau.t,sig.t,col=2,pch=16,cex=2)

# Smoothed distributions given posterior means for phi,tau2 and sig2
phi.m = mean(draws[ind,n+1])
tau.m = mean(draws[ind,n+2])
sig.m = mean(draws[ind,n+3])
kf1  = kalmanfilter(y,m0,C0,sig.m^2,phi.m,tau.m^2,10000)
qks1 = t(apply(kf1$x,1,quantile,c(0.05,0.5,0.95)))


# Unconditioal smoothed distributions
qmcmc = t(apply(draws[ind,1:n],2,quantile,c(0.05,0.5,0.95)))


par(mfrow=c(1,1))
ts.plot(qmcmc,lty=c(2,1,2),ylim=range(y,qmcmc))
lines(qks1[,1],col=2,lty=2)
lines(qks1[,2],col=2)
lines(qks1[,3],col=2,lty=2)
lines(qks[,1],col=4,lty=2)
lines(qks[,2],col=4)
lines(qks[,3],col=4,lty=2)
legend("topright",legend=c("KS (true theta)","KS (post mean)","KS (uncond.)"),col=c(1,2,4),lty=1)
points(y,pch=16)