#################################################################################################
sim.simplendlreg = function(n,mu,phi,tau2,sig2,beta0){
  tau   = sqrt(tau2)
  sig   = sqrt(sig2)
  beta  = rep(0,n)
  beta[1] = rnorm(1,mu+phi*beta0,tau)
  for (t in 2:n)
    beta[t] = rnorm(1,mu+phi*beta[t-1],tau) 
  x = rnorm(n)
  y = rnorm(n,beta*x,sig)
  return(list(x=x,y=y,beta=beta))
}

kalmanfiltersmoother = function(y,x,m0,c0,mu,phi,tau2,sig2){
  n    = length(y)
  m    = m0
  C    = C0
  phi2 = phi^2
  x2   = x^2
  a    = rep(0,n)
  R    = rep(0,n)
  mf   = rep(0,n)
  Cf   = rep(0,n)
  mb   = rep(0,n)
  Cb   = rep(0,n)
  for (t in 1:n){
    a[t]  = mu + phi*m
    R[t]  = phi2*C + tau2
    f     = x[t]*a[t]
    Q     = x2[t]*R[t] + sig2
    A     = R[t]*x[t]/Q
    m     = a[t] + A*(y[t]-f)
    C     = R[t] - A^2*Q
    mf[t] = m
    Cf[t] = C
  }
  mb[n] = mf[n]
  Cb[n] = Cf[n]
  for (t in (n-1):1){
    mb[t] = mf[t] + Cf[t]*phi/R[t+1]*(mb[t+1]-a[t+1])
    Cb[t] = Cf[t] - (Cf[t]^2*phi2)/(R[t+1]^2)*(R[t+1]-Cb[t+1])
  }
  return(list(mf=mf,Cf=Cf,mb=mb,Cb=Cb))
}

ffbs = function(y,x,m0,C0,mu,phi,tau2,sig2){
  n  = length(y)
  m  = m0
  C  = C0
  phi2 = phi^2
  x2   = x^2
  a  = rep(0,n)
  R  = rep(0,n)
  mf = rep(0,n)
  Cf = rep(0,n)
  for (t in 1:n){
    a[t]  = mu + phi*m
    R[t]  = phi2*C + tau2
    f     = x[t]*a[t]
    Q     = x2[t]*R[t] + sig2
    A     = R[t]*x[t]/Q
    m     = a[t]+A*(y[t]-f)
    C     = R[t] - A^2*Q
    mf[t] = m
    Cf[t] = C
  }
  beta    = rep(0,n)
  beta[n] = rnorm(1,mf[n],Cf[n])
  for (t in (n-1):1){
    var     = 1/(phi2/tau2+1/Cf[t])
    mean    = var*(beta[t+1]*phi/tau2+mf[t]/Cf[t])
    beta[t] = rnorm(1,mean,sqrt(var))  
  }
  return(beta)
}

simpledynamiclinearregression.MCMC = function(y,x,mu,phi,tau2,sig2,burnin,ndraw,thin){
  n      = length(y)
  niter  = burnin + thin*ndraw
  pars   = matrix(0,niter,4)
  betas  = matrix(0,niter,n)
  for (iter in 1:niter){
    beta   = ffbs(y,x,m0,C0,mu,phi,tau2,sig2)
    tau2   = 1/rgamma(1,(n-1)/2,sum((beta[2:n]-mu-phi*beta[1:(n-1)])^2)/2)
    sig2   = 1/rgamma(1,n/2,sum((y-x*beta)^2)/2)
    y1     = beta[2:n]-mu
    x1     = beta[1:(n-1)]
    var    = tau2/sum(x1^2)
    mean   = sum(x1*y1)/sum(x1^2)
    phi    = rnorm(1,mean,sqrt(var))
    y1     = beta[2:n]-phi*beta[1:(n-1)]
    mu     = rnorm(1,mean(y1),sqrt(tau2/(n-1)))
    pars[iter,]  = c(mu,phi,tau2,sig2)
    betas[iter,] = beta
  }
  ind = seq(burnin+1,niter,by=thin)
  pars  = pars[ind,]
  betas = betas[ind,]
  return(list(pars=pars,betas=betas))
}

#################################################################################################


# Simulating the data
set.seed(31416)

n     = 100
sig   = 1.0
tau   = 0.5
mu    = 0.0
phi   = 1.0
beta0 = 0.0

tau2  = tau^2
sig2  = sig^2
sim   = sim.simplendlreg(n,mu,phi,tau2,sig2,beta0)
x     = sim$x
y     = sim$y
beta  = sim$beta
param = c(mu,phi,tau2,sig2)

pdf(file="simple-normal-dynamic-linear-regression-FFBS-MCMC-bruteforceMCMC.pdf",width=12,height=8)
par(mfrow=c(1,2))
plot(beta,main="Time-varying regression coefficient",xlab="Time",
ylab=expression(beta))
for (i in seq(10,n,by=10))
  abline(v=i,lty=2)
for (i in 1:(n/10)){
  ind = (10*(i-1)+1):(10*i)
  points(10*(i-1)+5,mean(beta[ind]),pch=16,col=i,cex=2)
  lines(ind,beta[ind],col=i,lwd=2)
}

plot(x[1:10],y[1:10],xlim=range(x),ylim=range(y),xlab="X",ylab="Y")
abline(c(0,lm(y[1:10]~x[1:10]-1)$coef))
for (i in 2:(n/10)){
  ind = (10*(i-1)+1):(10*i)
  points(x[ind],y[ind],col=i)
  abline(c(0,lm(y[ind]~x[ind]-1)$coef),col=i)
}
abline(h=0,lty=2)
abline(v=0,lty=2)


##############################################################
# Kalman filter:   p(beta[t] | mu,phi,tau2,sig2,y[1],...,y[t])
# Kalman smoother: p(beta[t] | mu,phi,tau2,sig2,y[1],...,y[n])
##############################################################
m0  = 0
C0  = 10
run = kalmanfiltersmoother(y,x,m0,c0,mu,phi,tau2,sig2)
qf  = cbind(run$mf+qnorm(0.025)*sqrt(run$Cf),run$mf,run$mf+qnorm(0.975)*sqrt(run$Cf))
qb  = cbind(run$mb+qnorm(0.025)*sqrt(run$Cb),run$mb,run$mb+qnorm(0.975)*sqrt(run$Cb))


yrange = range(beta,qf,qb)
par(mfrow=c(1,2))
plot(beta,ylim=yrange,type="l",xlab="Time",ylab="beta",lwd=2)
title("Kalman filter")
lines(qf[,1],col=2)
lines(qf[,3],col=2)

plot(beta,ylim=yrange,type="l",xlab="Time",ylab="beta",lwd=2)
title("Kalman smoother")
lines(qb[,1],col=2)
lines(qb[,3],col=2)


#####################################
# Forward filtering backward sampling
# p(beta[t] | y[1],...,y[n])
#####################################
set.seed(3251241)
burnin   = 1000
ndraw    = 2000
thin     = 10
run.mcmc = simpledynamiclinearregression.MCMC(y,x,mu,phi,tau2,sig2,burnin,ndraw,thin)
pars     = run.mcmc$pars
betas    = run.mcmc$betas
qbeta    = t(apply(betas,2,quantile,c(0.025,0.5,0.975)))

names = c(expression(mu),expression(phi),expression(tau^2),expression(sigma^2))
par(mfrow=c(2,4))
for (i in 1:4){
  ts.plot(pars[,i],xlab="Iterations",ylab="",main=names[i])
  abline(h=param[i],col=2,lwd=2)
}
for (i in 1:4) 
  acf(pars[,i],main="")

par(mfrow=c(2,2))
for (i in 1:4){
  hist(pars[,i],prob=TRUE,main=names[i],xlab="")
  abline(v=param[i],col=2,lwd=2)
  abline(v=mean(pars[,i]),col=4,lwd=2)
}
legend("topright",legend=c("True value","Posterior mean"),col=c(2,4),lty=1,bty="n",lwd=2)


# Posterior summaries
qpars         = t(apply(pars,2,quantile,c(0.025,0.5,0.975)))
mpars         = apply(pars,2,mean)
tab           = cbind(param,mpars,qpars[,c(2,1,3)])
colnames(tab) = c("True","Mean","Median","2.5%","97.5%")
rownames(tab) = c("mu","phi","tau2","sig2")
round(tab,4)

#######################################################################
# Kalman smoother: p(beta[t] | mu,phi,tau2,sig2,y[1],...,y[n])
# (mu,phi,tau2,sig2) matching the posterior means and posterior medians
#######################################################################

run  = kalmanfiltersmoother(y,x,m0,c0,tab[1,2],tab[2,2],tab[3,2],tab[4,2])
qf1  = cbind(run$mf+qnorm(0.025)*sqrt(run$Cf),run$mf,run$mf+qnorm(0.975)*sqrt(run$Cf))
qb1  = cbind(run$mb+qnorm(0.025)*sqrt(run$Cb),run$mb,run$mb+qnorm(0.975)*sqrt(run$Cb))

run  = kalmanfiltersmoother(y,x,m0,c0,tab[1,3],tab[2,3],tab[3,3],tab[4,3])
qf2  = cbind(run$mf+qnorm(0.025)*sqrt(run$Cf),run$mf,run$mf+qnorm(0.975)*sqrt(run$Cf))
qb2  = cbind(run$mb+qnorm(0.025)*sqrt(run$Cb),run$mb,run$mb+qnorm(0.975)*sqrt(run$Cb))


##############################################################
# Comparison
##############################################################
par(mfrow=c(1,2))
ts.plot(qf[,c(1,3)],ylim=range(qf,qf1),ylab=expression(beta),main="")
title("Marginal filtered distributions of beta[t]\n Given y[1],...,y[t]")
lines(qf1[,1],col=4)
lines(qf1[,3],col=4)
lines(qf2[,1],col=6)
lines(qf2[,3],col=6)

legend("topright",legend=c("Given true (mu,phi,tau2,sig2)",
"Given mean of (mu,phi,tau2,sig2)",
"Given median of (mu,phi,tau2,sig2)"),col=c(1,4,6),bty="n",lty=1)

ts.plot(qbeta[,c(1,3)],ylim=range(qb,qb1,qbeta),ylab=expression(beta),main="")
title("Marginal smoothed distributions of beta[t]\n Given y[1],...,y[n]")
lines(qb[,1],col=2,lwd=3)
lines(qb[,3],col=2,lwd=3)
lines(qb1[,1],col=4)
lines(qb1[,3],col=4)
lines(qb2[,1],col=6)
lines(qb2[,3],col=6)
legend("topright",legend=c("Given true (mu,phi,tau2,sig2)",
"Given mean of (mu,phi,tau2,sig2)",
"Given median of (mu,phi,tau2,sig2)",
"Integrating out (mu,phi,tau2,sig2)"),col=c(1,4,6,2),bty="n",lty=1,lwd=c(1,1,1,3))









set.seed(3251241)
burnin = 1000
ndraw  = 1000
thin   = 1
ts     = seq(50,n,by=1)
nt     = length(ts)
qpars  = array(0,c(nt,4,3))
qbetas = matrix(0,nt,3)
for (i in 1:nt){
  ind = 1:ts[i]
  y1 = y[ind]
  x1 = x[ind]
  run.mcmc   = simpledynamiclinearregression.MCMC(y1,x1,mu,phi,tau2,sig2,burnin,ndraw,thin)
  qpars[i,,] = t(apply(run.mcmc$pars,2,quantile,c(0.025,0.5,0.975)))
  qbetas[i,] = quantile(run.mcmc$betas[,ts[i]],c(0.025,0.5,0.975))
}

par(mfrow=c(2,2))
for (i in 1:4){
  plot(ts,qpars[,i,2],ylim=range(qpars[,i,]),ylab="",type="l",main=names[i],xlab="Time")
  lines(ts,qpars[,i,1],col=2)
  lines(ts,qpars[,i,3],col=2)
}

par(mfrow=c(1,1))
plot(ts,qbetas[,1],ylim=range(qbetas),ylab="",type="l",main=expression(beta),xlab="Time")
lines(ts,qbetas[,3])
lines(ts,qf1[ts,1],col=3)
lines(ts,qf1[ts,3],col=3)
lines(ts,qf2[ts,1],col=4)
lines(ts,qf2[ts,3],col=4)

dev.off()