#######################################################
#
# 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).
#
# Sampling x1,...,xn jointly (FFBS)
# Sampling x1,...,xn individually
#
#######################################################

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

singlemove = function(y,sig2,phi,tau2,m0,C0,x){
  n    = length(y)
  var  = 1/(phi^2/tau2+1/(phi^2*C0+tau2)+1/sig2)
  mean = var*(phi*m0/(phi^2*C0+tau2)+phi*x[2]/tau2+y[1]/sig2)
  x[1] = rnorm(1,mean,sqrt(var))
  for (t in 2:(n-1)){
    var = 1/((1+phi^2)/tau2+1/sig2)
    mean = var*(phi*(x[t-1]+x[t+1])/tau2+y[t]/sig2)
    x[t] = rnorm(1,mean,sqrt(var))  	
  }
  var = 1/(1/tau2+1/sig2)
  mean = var*(phi*x[n-1]/tau2+y[n]/sig2)
  x[n] = rnorm(1,mean,sqrt(var))
  return(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.t*x[t-1]+rnorm(1,0,tau.t)
x.t = x[(n+1):(2*n)]
y  = rnorm(n,x.t,sig.t)

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


#######################################################
# 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)
# x0~N(m0,C0)
f0=0
F0=100
nu0=1
tau20=1
eta0=1
sig20=1
m0 = 0
C0 = 1
 
#######################################################
# Jointly sampling x1,...xn from p(x1,...,xn|y1,...,yn)
#######################################################

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

# MCMC setup
M0    = 1000
M     = 1000
L     = 1
niter = M0+M*L
draws = matrix(0,niter,n+3)
system.time({
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))
}
})
draws1 = draws

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,1))
acf = acf(draws[ind,1],plot=FALSE)
plot(acf$acf,type="l",col=grey(0.7),xlab="LAG",ylab="ACF",ylim=c(-0.2,1))
for (t in 2:n){
  acf = acf(draws[ind,t],plot=FALSE)
  lines(acf$acf,col=grey(0.7))
}
abline(h=0,lwd=2)

# Unconditioal smoothed distributions
qmcmc1 = t(apply(draws[ind,1:n],2,quantile,c(0.05,0.5,0.95)))
par(mfrow=c(1,1))
ts.plot(qmcmc1,lty=c(2,1,2),ylim=range(y,qmcmc1))


#######################################################
# Sampling x1,...xn from full conditionals:
# p(xt|x1,...,xt-1,xt+1,...,xn,y1,...,yn)
#######################################################

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

x = x.t
phi=phi.t
tau2 = tau.t^2
sig2 = sig.t^2

# MCMC setup
M0    = 0
M     = 1000
L     = 1
niter = M0+M*L
draws = matrix(0,niter,n+3)
system.time({
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 = singlemove(y,sig2,phi,tau2,m0,C0,x)
  draws[iter,]=c(x,phi,sqrt(tau2),sqrt(sig2))
}
})
draws2 = draws

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,1))
acf = acf(draws[ind,1],plot=FALSE)
plot(acf$acf,type="l",col=grey(0.7),xlab="LAG",ylab="ACF",ylim=c(-0.2,1))
for (t in 2:n){
  acf = acf(draws[ind,t],plot=FALSE)
  lines(acf$acf,col=grey(0.7))
}
abline(h=0,lwd=2)

# Unconditioal smoothed distributions
qmcmc2 = t(apply(draws[ind,1:n],2,quantile,c(0.05,0.5,0.95)))
par(mfrow=c(1,1))
ts.plot(qmcmc2,lty=c(2,1,2),ylim=range(y,qmcmc2))





par(mfrow=c(1,1))
acf = acf(draws1[ind,1],plot=FALSE)
plot(acf$acf,type="l",col=grey(0.5),xlab="LAG",ylab="ACF",ylim=c(-0.2,1))
for (t in 2:n){
  acf = acf(draws1[ind,t],plot=FALSE)
  lines(acf$acf,col=grey(0.5))
}
for (t in 1:n){
  acf = acf(draws2[ind,t],plot=FALSE)
  lines(acf$acf,col=2)
}

abline(h=0,lwd=2)
legend("topright",legend=c("Joint move","Single move"),col=c(grey(0.5),2),lty=1,lwd=3)

par(mfrow=c(1,1))
ts.plot(qmcmc1,lty=c(2,1,2),ylim=range(qmcmc1,qmcmc2))
lines(qmcmc2[,1],col=2,lty=2)
lines(qmcmc2[,2],col=2)
lines(qmcmc2[,3],col=2,lty=2)