set.seed(12345)
n      = 100
sig    = 1
tau    = 0.25
alpha  = 0.1
beta   = 0.9
theta0 = 0
sig2   = sig^2
tau2   = tau^2
true   = c(alpha,beta,tau2,sig2)

theta    = rep(0,n)
theta[1] = rnorm(1,alpha+beta*theta0,tau)
for (t in 2:n)
  theta[t] = rnorm(1,alpha+beta*theta[t-1],tau)
y = rnorm(n,theta,sig)

plot(y)
lines(theta,col=2)


############################################################
# Learning the latent variables theta[1],...,theta[n]
############################################################
# theta[1] ~ N(a1,R1)
a1=0
R1=1

# a few constants
var1 = 1/(1/sig2+1/tau2+1/R1)
vart = 1/(1/sig2+2/tau2)
varn = 1/(1/sig2+1/tau2)


# initial value
thetas = y

# MCMC setup
burnin = 2000
M = 2000
step = 1
niter = burnin+M*step

par(mfrow=c(1,1))
plot(y)
thetas.draw = matrix(0,niter,n)
for (iter in 1:niter){
  # t=1
  mean = var1*(y[1]/sig2+thetas[2]/tau2+a1/R1)
  thetas[1] = rnorm(1,mean,sqrt(var1))

  # t=n
  mean = varn*(y[n]/sig2+thetas[n-1]/tau2)
  thetas[n] = rnorm(1,mean,sqrt(varn))

  # t=2,...,n-1
  for (t in 2:(n-1)){
    mean = vart*(y[t]/sig2+(thetas[t-1]+thetas[t+1])/tau2)
    thetas[t] = rnorm(1,mean,sqrt(vart))
  }
  lines(thetas,col=3)
  thetas.draw[iter,] = thetas
}
points(y,pch=16)
lines(theta,col=2,lwd=2)

thetas.draw = thetas.draw[(burnin+1):niter,]

par(mfrow=c(3,5))
few = sort(sample(1:n,size=5))
for (i in 1:5) 
  ts.plot(thetas.draw[,few[i]],xlab="Iterations",ylab="",
  main=paste("theta[",few[i],"]",sep=""))
for (i in 1:5) acf(thetas.draw[,few[i]],main="")
for (i in 1:5) hist(thetas.draw[,few[i]],prob=TRUE,main="",xlab="")


par(mfrow=c(1,1))
k = 40
plot(1:k,acf(thetas.draw[1,],lag.max=k,plot=F)$acf[2:(k+1)],
     ylab="ACF",type="l",xlab="Lag",ylim=c(-0.1,1))
for (i in 2:M)
  lines(acf(thetas.draw[i,],lag.max=k,plot=FALSE)$acf[2:(k+1)],
        col=grey(0.9))
abline(h=0,lty=2)  
  
qthetas = t(apply(thetas.draw,2,quantile,c(0.025,0.5,0.975)))

par(mfrow=c(1,1))
plot(y,xlab="Time",pch=16)
lines(theta,col=2,lwd=2)
lines(qthetas[,1],col=4,lwd=2)
lines(qthetas[,2],col=4,lwd=2)
lines(qthetas[,3],col=4,lwd=2)

corr = NULL
for (i in 2:n)
  for (j in 1:(i-1))
     corr = c(corr,cor(thetas.draw[,c(i,j)])[1,2])

library(reshape2)
library(ggplot2)
corr1 <- melt(cor(thetas.draw))
ggplot(data = corr1, aes(x=Var1, y=Var2, fill=value)) + 
  geom_tile()
  



############################################################
# Learning the constant parameters as well
############################################################
# prior hyperparameters
a.sig = 1
b.sig = 1
a.tau = 1
b.tau = 1
d0    = c(0,0)
B0    = diag(9,2)
iD0   = solve(B0)
iD0d0 = iD0%*%d0

# initial value
alpha  = true[1]
beta   = true[2]
tau2   = true[3]
sig2   = true[4]
thetas = y

# MCMC setup
burnin = 2000
M      = 2000
niter  = burnin+M
thetas.draw1 = matrix(0,niter,n+4)
for (iter in 1:niter){
  # t=1
  C1 = 1/(1/sig2+1/R1+beta^2/tau2)
  m1 = C1*(y[1]/sig2+a1/R1+beta*(thetas[2]-alpha)/tau2)
  thetas[1] = rnorm(1,m1,sqrt(C1))

  # t=2,...,n-1
  for (t in 2:(n-1)){
  	Ct = 1/(1/sig2+(1+beta^2)/tau2)
    mt = Ct*(y[t]/sig2+(alpha*(1-beta)+
         beta*(thetas[t-1]+thetas[t+1]))/tau2)
    thetas[t] = rnorm(1,mt,sqrt(Ct))
  }
  
  # t=n
  Cn = 1/(1/sig2+1/tau2)
  mn = Cn*(y[n]/sig2+(alpha+beta*thetas[n-1])/tau2)
  thetas[n] = rnorm(1,mn,sqrt(Cn))

  # sig2
  par1 = a.sig+n/2
  par2 = b.sig+sum((y-theta)^2)/2
  sig2 = 1/rgamma(1,par1,par2)
  
  # tau2
  par1 = a.tau+(n-1)/2
  par2 = b.tau+sum((theta[2:n]-alpha-beta*theta[1:(n-1)])^2)/2
  tau2 = 1/rgamma(1,par1,par2)

  # (alpha,beta)
  X     = cbind(1,theta[1:(n-1)])
  yt    = theta[2:n]
  var   = solve(t(X)%*%X/tau2+iD0)
  mean  = var%*%(t(X)%*%yt/tau2+iD0d0) 
  draw  = mean+t(chol(var))%*%rnorm(2) 
  alpha = draw[1]
  beta  = draw[2]
  thetas.draw1[iter,] = c(thetas,alpha,beta,tau2,sig2)
}
thetas.draw1 = thetas.draw1[(burnin+1):niter,]

qthetas1 = t(apply(thetas.draw1[,1:n],2,quantile,c(0.025,0.5,0.975)))


names = c("alpha","beta","tau2","sig2")
par(mfrow=c(3,4))
for (i in 1:4)
  ts.plot(thetas.draw1[,n+i],ylab="",xlab="Iterations",main=names[i])
for (i in 1:4)
  acf(thetas.draw1[,n+i],ylab="",main="")
for (i in 1:4){
  hist(thetas.draw1[,n+i],prob=TRUE,xlab="",main="")
  abline(v=true[i],col=2,lwd=2)  
}

par(mfrow=c(1,1))
plot(thetas.draw1[,(n+1):(n+2)],xlab=expression(alpha),
     ylab=expression(beta),main="Posterior draws")

par(mfrow=c(1,1))
plot(y,xlab="Time",pch=16)
lines(qthetas[,1],col=4,lwd=2)
lines(qthetas[,2],col=4,lwd=2)
lines(qthetas[,3],col=4,lwd=2)
lines(qthetas1[,1],col=6,lwd=2)
lines(qthetas1[,2],col=6,lwd=2)
lines(qthetas1[,3],col=6,lwd=2)
legend("topleft",legend=c("Learning latent states","Learning states and constant parameters"),col=c(4,6),lwd=2,bty="n",lty=1)
  
  
######################################################
# Now, let us implement FFBS                         #
# Forward filtering, backward sampling               #
#                                                    #
#  y(t)     ~ N(alpha(t)+F(t)*theta(t);V(t))         #
#  theta(t) ~ N(gamma+G*theta(t-1);W)                #
#                                                    #
######################################################
ffbs = function(y,F,alpha,V,G,gamma,W,a1,R1,nd=1){
  n = length(y)
  if (length(F)==1){F = rep(F,n)}
  if (length(alpha)==1){alpha = rep(alpha,n)}
  if (length(V)==1){V = rep(V,n)}
  a = rep(0,n)
  R = rep(0,n)
  m = rep(0,n)
  C = rep(0,n)
  B = rep(0,n-1)
  H = rep(0,n-1)
  # time t=1
  a[1] = a1
  R[1] = R1
  f    = alpha[1]+F[1]*a[1]
  Q    = R[1]*F[1]**2+V[1]
  A    = R[1]*F[1]/Q
  m[1] = a[1]+A*(y[1]-f)
  C[1] = R[1]-Q*A**2
  # forward filtering
  for (t in 2:n){
    a[t] = gamma + G*m[t-1]
    R[t] = C[t-1]*G**2 + W
    f    = alpha[t]+F[t]*a[t]
    Q    = R[t]*F[t]**2+V[t]
    A    = R[t]*F[t]/Q
    m[t] = a[t]+A*(y[t]-f)
    C[t] = R[t]-Q*A**2
    B[t-1] = C[t-1]*G/R[t]
    H[t-1] = sqrt(C[t-1]-R[t]*B[t-1]**2)
  }
  # backward sampling
  theta = matrix(0,nd,n)
  theta[,n] = rnorm(nd,m[n],sqrt(C[n]))
  for (t in (n-1):1)
    theta[,t] = rnorm(nd,m[t]+B[t]*(theta[,t+1]-a[t+1]),H[t])
  if (nd==1){
    theta[1,]
  }
  else{
    theta
  }
}


# initial value
alpha  = true[1]
beta   = true[2]
tau2   = true[3]
sig2   = true[4]
thetas = y

# MCMC setup
burnin = 100
M      = 1000
niter  = burnin+M
thetas.draw2 = matrix(0,niter,n+4)
mf = rep(0,n)
Cf = rep(0,n)
for (iter in 1:niter){
  # theta1,...,thetan jointly
  thetas = ffbs(y,1,0,sig2,beta,alpha,tau2,a1,R1,nd=1)

# #   A = R1/(R1+sig2)
  # mf[1] = (1-A)*a1+A*y[1]
  # Cf[1] = (1-A)*R1
  # for (t in 2:n){
    # a = alpha+beta*mf[t-1]
    # R = beta^2*Cf[t-1]
    # A = R/(R+sig2)
    # mf[t] = (1-A)*a+A*y[t]
    # Cf[t] = (1-A)*R 
  # }
  # mb=mf[n]
  # Cb=Cf[n]
  # thetas[n] = rnorm(1,mb,sqrt(Cb))
  # for (t in (n-1):1){
  	# Cb = 1/(1/Cf[t]+beta^2/tau2)
  	# mb = Cb*(mf[t]/Cf[t]+beta*(thetas[t+1]-alpha)/tau2)
    # thetas[t] = rnorm(1,mb,sqrt(Cb))
  # }

  # sig2
  par1 = a.sig+n/2
  par2 = b.sig+sum((y-theta)^2)/2
  sig2 = 1/rgamma(1,par1,par2)
  
  # tau2
  par1 = a.tau+(n-1)/2
  par2 = b.tau+sum((theta[2:n]-alpha-beta*theta[1:(n-1)])^2)/2
  tau2 = 1/rgamma(1,par1,par2)

  # (alpha,beta)
  X     = cbind(1,theta[1:(n-1)])
  yt    = theta[2:n]
  var   = solve(t(X)%*%X/tau2+iD0)
  mean  = var%*%(t(X)%*%yt/tau2+iD0d0) 
  draw  = mean+t(chol(var))%*%rnorm(2) 
  alpha = draw[1]
  beta  = draw[2]
  thetas.draw2[iter,] = c(thetas,alpha,beta,tau2,sig2)
}
thetas.draw2 = thetas.draw2[(burnin+1):niter,]

qthetas2 = t(apply(thetas.draw2[,1:n],2,quantile,c(0.025,0.5,0.975)))


names = c("alpha","beta","tau2","sig2")
par(mfrow=c(3,4))
for (i in 1:4)
  ts.plot(thetas.draw2[,n+i],ylab="",xlab="Iterations",main=names[i])
for (i in 1:4)
  acf(thetas.draw2[,n+i],ylab="",main="")
for (i in 1:4){
  hist(thetas.draw2[,n+i],prob=TRUE,xlab="",main="")
  abline(v=true[i],col=2,lwd=2)  
}

par(mfrow=c(1,1))
plot(thetas.draw2[,(n+1):(n+2)],xlab="alpha",ylab="beta",main="Posterior draws")

par(mfrow=c(1,2))
k = 15
plot(1:k,acf(thetas.draw[1,],lag.max=k,plot=F)$acf[2:(k+1)],
     ylab="ACF",type="l",xlab="Lag",ylim=c(-0.5,1))
for (i in 2:M)
  lines(acf(thetas.draw[i,],lag.max=k,plot=FALSE)$acf[2:(k+1)],
        col=grey(0.9))
abline(h=0,lty=2)  
title("Sampling theta(t)|theta(-t)")

plot(1:k,acf(thetas.draw2[1,],lag.max=k,plot=F)$acf[2:(k+1)],
     ylab="ACF",type="l",xlab="Lag",ylim=c(-0.5,1))
for (i in 2:M)
  lines(acf(thetas.draw2[i,],lag.max=k,plot=FALSE)$acf[2:(k+1)],
        col=grey(0.9))
abline(h=0,lty=2)  
title("Sampling (theta(1),...,theta(n)) jointly\nFFBS")


par(mfrow=c(1,1))
plot(y,xlab="Time",pch=16)
lines(qthetas[,1],col=4,lwd=2)
lines(qthetas[,2],col=4,lwd=2)
lines(qthetas[,3],col=4,lwd=2)
lines(qthetas1[,1],col=6,lwd=2)
lines(qthetas1[,2],col=6,lwd=2)
lines(qthetas1[,3],col=6,lwd=2)
lines(qthetas2[,1],col=3,lwd=2)
lines(qthetas2[,2],col=3,lwd=2)
lines(qthetas2[,3],col=3,lwd=2)
legend("topleft",legend=c("Learning latent states","Learning states and constant parameters"),col=c(4,6),lwd=2,bty="n",lty=1)