##################################################################################
#
# Fitting a simple dynamic regression model
#
###################################################################################

############################################################
# 1st case
############################################################

# Simulating the data
# -------------------
set.seed(12345)
n     = 300
beta1 = 4
beta2 = 1
beta3 = -1
sig2  = 4
x     = rnorm(n)
y     = rep(0,n)
sig   = sqrt(sig2)
y[1:(n/3)]   = beta1*x[1:(n/3)]   + rnorm(n/3,0,sig)
y[(n/3+1):(2*n/3)] = beta2*x[(n/3+1):(2*n/3)] + rnorm(n/3,0,sig)
y[(2*n/3+1):n] = beta3*x[(2*n/3+1):n] + rnorm(n/3,0,sig)

par(mfrow=c(2,2))
plot(x,y,col=c(rep(1,n/3),rep(2,n/3),rep(3,n/3)),pch=16)
plot(x,y,pch=16)
ts.plot(y)
ts.plot(x)

# Prior hyperparameters
m0   = 0
C0   = 1
a0   = 5
b0   = 5
c0   = 5
d0   = 5

# MCMC scheme
sig2  = 1
tau2  = 1
M0    = 1000
M     = 1000
niter = M0+M
mf    = rep(0,n)
Cf    = rep(0,n)
draws = matrix(0,niter,n+2)
m     = m0
C     = C0
for (iter in 1:niter){
  # Sampling the states jointly
  # Forward filtering, backward sampling (FFBS)
  for (t in 1:n){
    at  = m
    Rt  = C+tau2
    ft    = x[t]*at
    Qt    = x[t]^2*Rt+sig2
    At    = Rt*x[t]/Qt
    m     = at + At*(y[t]-ft)
    C     = Rt - At^2*Qt
    mf[t] = m
    Cf[t] = C
  }
  betas = rep(0,n)
  betas[n] = rnorm(1,mf[n],sqrt(Cf[n]))
  for (t in (n-1):1){
    var = 1/(1/tau2+1/Cf[t])
    mean = var*(betas[t+1]/tau2+mf[t]/Cf[t])
    betas[t] = rnorm(1,mean,sqrt(var))	
  }

  # sampling sig2
  sig2 = 1/rgamma(1,a0+n/2,b0+sum((y-x*betas)^2)/2)

  # sampling tau2
  tau2 = 1/rgamma(1,a0+(n-1)/2,b0+sum((betas[2:n]-betas[1:(n-1)])^2)/2)

  # storing draws
  draws[iter,] = c(betas,sqrt(sig2),sqrt(tau2))
}
draws = draws[(M0+1):niter,]
qb = t(apply(draws[,1:n],2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(2,3))
ts.plot(draws[,n+1],main=expression(sigma),ylab="")
acf(draws[,n+1],main="")
hist(draws[,n+1],prob=TRUE,main="",xlab="")
ts.plot(draws[,n+2],main=expression(tau),ylab="")
acf(draws[,n+2],main="")
hist(draws[,n+2],prob=TRUE,main="",xlab="")

par(mfrow=c(1,1))
ts.plot(qb,col=c(3,2,3))
segments(0,beta1,n/3,beta1,col=4,lwd=2)
segments(n/3,beta2,2*n/3,beta2,col=4,lwd=2)
segments(2*n/3,beta3,n,beta3,col=4,lwd=2)
abline(v=n/3,col=4,lty=2)
abline(v=2*n/3,col=4,lty=2)




############################################################
# 2nd case
############################################################


# Simulating the data
# -------------------
set.seed(654321)
n     = 300
sig2  = 4
tau2  = 0.1
x     = rnorm(n)
beta  = rep(0,n)
for (t in 2:n)
  beta[t] = beta[t-1]+rnorm(1,0,sqrt(tau2))
y = beta*x+rnorm(n,0,sqrt(sig2))

par(mfrow=c(1,2))
ts.plot(beta)
plot(x,y)

# Prior hyperparameters
m0   = 0
C0   = 1
a0   = 5
b0   = 5
c0   = 5
d0   = 5

# MCMC scheme
sig2  = 1
tau2  = 1
M0    = 1000
M     = 1000
niter = M0+M
mf    = rep(0,n)
Cf    = rep(0,n)
draws = matrix(0,niter,n+2)
m     = m0
C     = C0
for (iter in 1:niter){
  # Sampling the states jointly
  # Forward filtering, backward sampling (FFBS)
  for (t in 1:n){
    at  = m
    Rt  = C+tau2
    ft    = x[t]*at
    Qt    = x[t]^2*Rt+sig2
    At    = Rt*x[t]/Qt
    m     = at + At*(y[t]-ft)
    C     = Rt - At^2*Qt
    mf[t] = m
    Cf[t] = C
  }
  betas = rep(0,n)
  betas[n] = rnorm(1,mf[n],sqrt(Cf[n]))
  for (t in (n-1):1){
    var = 1/(1/tau2+1/Cf[t])
    mean = var*(betas[t+1]/tau2+mf[t]/Cf[t])
    betas[t] = rnorm(1,mean,sqrt(var))	
  }

  # sampling sig2
  sig2 = 1/rgamma(1,a0+n/2,b0+sum((y-x*betas)^2)/2)

  # sampling tau2
  tau2 = 1/rgamma(1,a0+(n-1)/2,b0+sum((betas[2:n]-betas[1:(n-1)])^2)/2)

  # storing draws
  draws[iter,] = c(betas,sqrt(sig2),sqrt(tau2))
}
draws = draws[(M0+1):niter,]
qb = t(apply(draws[,1:n],2,quantile,c(0.05,0.5,0.95)))

par(mfrow=c(2,3))
ts.plot(draws[,n+1],main=expression(sigma),ylab="")
acf(draws[,n+1],main="")
hist(draws[,n+1],prob=TRUE,main="",xlab="")
ts.plot(draws[,n+2],main=expression(tau),ylab="")
acf(draws[,n+2],main="")
hist(draws[,n+2],prob=TRUE,main="",xlab="")

par(mfrow=c(1,1))
ts.plot(qb,lty=c(2,1,2))
lines(beta,col=2,lwd=3)
legend("topright",legend=c("True state","Posterior"),lty=1,col=2:1,lwd=2)